advent-of-code/2020/day7/day7.py

#!/usr/bin/env python3
import re
from collections import defaultdict, deque


def main(inp):
    with open(inp) as input_rules:
        rules = parse_rules(input_rules)
        reverse_rules = build_reverse_rules(rules)
        print(part1(reverse_rules))
        print(part2(rules, "shiny gold"))


def parse_rules(input_rules):
    rules = {}
    for input_rule in input_rules:
        color, rule = input_rule.split(" bags contain ")
        rules[color] = {color: int(number) for number, color in re.findall('(\d+) (\w+ \w+)', rule)}
    return rules


def build_reverse_rules(rules):
    reverse_rules = defaultdict(list)
    for bag, inner_rules in rules.items():
        for c in inner_rules:
            reverse_rules[c].append(bag)
    return reverse_rules


def part1(reverse_rules):
    queue  = deque(("shiny gold",))
    may_contain_shiny_gold = set()
    while queue:
        color = queue.pop()
        for c in reverse_rules.get(color, []):
            if c not in may_contain_shiny_gold:
                may_contain_shiny_gold.add(c)
                queue.appendleft(c)
    return len(may_contain_shiny_gold)


def part2(rules, color):
    return sum(number + number * part2(rules, c) for c, number in rules[color].items())


if __name__ == "__main__":
    main("input.txt")