advent03.py

from typing import Collection


def count_bits(diagnostics: Collection[str]) -> (list[int], int):
    bit_counts = []
    for number in diagnostics:
        for bit_index, binary_digit in enumerate(number):
            while bit_index >= len(bit_counts):
                bit_counts.append(0)
            
            bit_counts[bit_index] += int(binary_digit, 2)
            
    return list(bit_counts), len(diagnostics) / 2


def filter_ratings(candidates: Collection[str], predicate: callable) -> int:
    for bit_index in range(bit_width):
        counts, majority = count_bits(candidates)
        candidates = {n for n in candidates if int(n[bit_index], 2) == predicate(counts[bit_index], majority)}

        if len(candidates) <= 1:
            break

    rating ,= candidates
    return int(rating, 2)


def calc_o2_co2(diagnostics: Collection[str]) -> (int, int):
    return filter_ratings(diagnostics, lambda v, m: v >= m), filter_ratings(diagnostics, lambda v, m: v < m)


def calc_gamma_epsilon(diagnostics: Collection[str]) -> (int, int):
    bit_counts, majority = count_bits(diagnostics)
    gamma, epsilon = 0, 0
    # We reverse the bit counts to associate their digit index with their power of two.
    for bit_index, bit_count in enumerate(reversed(bit_counts)):
        if bit_count >= majority:
            gamma += 2 ** bit_index
        if bit_count <= majority:
            epsilon += 2 ** bit_index
    return gamma, epsilon


if __name__ == '__main__':
    diagnostic_data = [line.strip() for line in open("diagnostics.txt").readlines()]

    c, e, = calc_gamma_epsilon(diagnostic_data)
    print(f"Power consumption at {c * e}")

    oxy, co2, = calc_o2_co2(diagnostic_data)
    print(f"Life support rating {oxy * co2}")