day7.py

# https://adventofcode.com/2020/day/7
from pathlib import Path
import re

INPUT_FILE = Path(__file__).parent / "input.txt"

NEEDLE = "shiny gold"

def parse_rules(stream):
    rules = {}
    for line in stream:
        bag = re.match("[a-z]+ [a-z]+", line).group()
        contents = re.findall("(\d+) ([a-z]+ [a-z]+)", line)
        rules[bag] = {name: int(amount) for amount, name in contents}
    return rules

def get_super_bags(rules, needle):
    queue = set([needle])
    result = set()
    while queue:
        needle = queue.pop()
        for bag, contents in rules.items():
            if needle in contents:
                queue.add(bag)
                result.add(bag)
    return result

def count_total_bags(rules, needle):
    sub_amounts = (
        amount * count_total_bags(rules, bag)
        for bag, amount in rules[needle].items()
    )
    return 1 + sum(sub_amounts)

def main():
    with INPUT_FILE.open() as stream:
        rules = parse_rules(stream)
    print(len(get_super_bags(rules, NEEDLE)))
    print(count_total_bags(rules, NEEDLE) - 1)