advent_3_part2.py

COUNT = 0
DATA = 1

class Trie:
  def __init__(self):
    self.root = Node()
  def add_string(self, string):
    node = self.root
    for char in string:
      node = node.add(char)

class Node:
  def __init__(self):
    self.childs = {}

  def add(self, char):
    if not char in self.childs:
      self.childs[char] = [0, Node()]
    self.childs[char][COUNT] += 1
    return self.childs[char][DATA]

def get_rating(node, bits_length, best_fn, keep_on_tie):
  rating = 0
  position = 0
  while True:
    if not node.childs:
      return rating
    best_node = None
    best_count = None
    best_char = None
    for char in node.childs:
      count, child = node.childs[char]
      if best_node is None or best_fn(count, best_count) or (count == best_count and char == keep_on_tie):
        best_node = child
        best_count = count
        best_char = char
    node = best_node
    if best_char == '1':
      rating += 2 ** (bits_length - 1 - position)
    position += 1

if __name__ == '__main__':
  with open('input.txt') as data:
    trie = Trie()
    bits_length = 0
    for line in data:
      line = line.strip()
      trie.add_string(line)
      bits_length = len(line)
    oxygen = get_rating(trie.root, bits_length, lambda x, y: x > y, '1')
    co2 = get_rating(trie.root, bits_length, lambda x, y: x < y, '0')
    print(oxygen, co2, oxygen * co2)