aoc-2021/3.py

from util import get_input

input = get_input('3.input')

def count(input):
    counts = [0 for _ in input[0]]
    for line in input:
        for i, c in enumerate(line):
            counts[i] += 1 if c == '1' else 0
    return counts

gamma = ''.join(['1' if c > len(input) / 2 else '0' for c in count(input)])
epsilon = ''.join(['0' if c > len(input) / 2 else '1' for c in count(input)])

print(int(gamma, 2) * int(epsilon, 2))

def search(input, invert):
    remaining = list(input)
    pos = 0
    while len(remaining) > 1:
        print(len(remaining))
        cmp = count(remaining)[pos] >= len(remaining) / 2
        if invert:
            cmp = not cmp
        target = '1' if cmp else '0'
        remaining = [a for a in remaining if a[pos] == target]
        pos += 1
    return remaining[0]

ox = search(input, False)
co = search(input, True)

print(int(ox, 2) * int(co, 2))