from typing import Generator from collections.abc import Iterable from collections import Counter Bits = tuple[bool, ...] Input = list[Bits] def column(rows: Iterable[Bits], n: int) -> Generator[bool, None, None]: for row in rows: yield row[n] def most_common(bits: Iterable[bool]) -> bool: freqs: list[tuple[bool, int]] = Counter(bits).most_common() if freqs[0][1] == freqs[1][1]: return True return freqs[0][0] def bits_to_int(bits: Bits) -> int: ret: int = 0 for i, b in enumerate(reversed(bits)): if b: ret += 2 ** i return ret def part1(nums: Input) -> int: nbits: int = len(nums[0]) gamma: Bits = tuple(most_common(column(inp, i)) for i in range(nbits)) epsilon: Bits = tuple(not bit for bit in gamma) return bits_to_int(gamma) * bits_to_int(epsilon) def part2_impl(nums: Input, complement: bool, bit: int = 0) -> Bits: if len(nums) == 1: return nums[0] target: bool = most_common(column(nums, bit)) if complement: target = not target nums = list(filter(lambda n: n[bit] == target, nums)) return part2_impl(nums, complement, bit + 1) def part2(nums: Input) -> int: oxygen: Bits = part2_impl(nums, False) co2: Bits = part2_impl(nums, True) return bits_to_int(oxygen) * bits_to_int(co2) inp: Input = [tuple(bool(int(c)) for c in l.strip()) for l in open('3.in')] print(part1(inp)) print(part2(inp))