#!/usr/bin/env python
# coding: utf-8

# In[38]:


def count(f):
    sums = {}
    lines = 0
    for l in f:
        for p, ch in enumerate(l.rstrip()):
            sums[p] = sums.get(p, 0) + int(ch)
        lines += 1

    return sums, lines

def rates(f):
    sums, lines = count(f)
    digits = len(sums)
    gamma = 0
    epsilon = 0
    for p, s in sums.items():
        if s > lines/2:
            gamma |= 1<<(digits-p-1)
        else:
            epsilon |= 1<<(digits-p-1)

    return gamma, epsilon

import import_ipynb
import helper
f = helper.open_file('2021_3.txt')

gamma, epsilon = rates(f)

print(gamma*epsilon)


# In[39]:


class Node:
    def __init__(self):
        self.children = [None, None]
        
def insert(node, bitstring):
    for bit in bitstring:
        bit = int(bit)
        if node.children[bit] is None:
            node.children[bit] = Node()
        node = node.children[bit]
        
def count_leafs(node):
    if node is None:
        return 0
    
    c = 0
    
    leaf = True
    for n in node.children:
        if n:
            c += count_leafs(n)
            leaf = False
            
    if leaf:
        c += 1
    
    return c

import import_ipynb
import helper
f = helper.open_file('2021_3.txt')

root = Node()

for l in f:
    insert(root, l.rstrip())
    
def rating(node, bit_criteria):
    r = 0
    while node:        
        c0 = count_leafs(node.children[0])
        c1 = count_leafs(node.children[1])
        
        if c0 == 0 and c1 == 0:
            break
    
        r <<= 1
        if not (c0 == 0 and c1 == 1) and \
            (bit_criteria(c0, c1) or (c0 == 1 and c1 == 0)):
            
            node = node.children[0]
        else:
            node = node.children[1]
            r |= 1          
            
    return r

import operator

oxygen_generator_rate = rating(root, operator.gt)
c02_scrubber_rate =  rating(root, operator.le)

oxygen_generator_rate * c02_scrubber_rate


# In[ ]: