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

#!/usr/bin/env python

from collections import defaultdict
import regex
import sys

RULE_REGEX = regex.compile(r'(.*) bags contain (?:(no|\d+) (.*?) bags?(?:, )?)+\.')

def process_rules(rules):
    rules_dict = {}
    for rule in rules:
        m = regex.fullmatch(RULE_REGEX, rule)
        if m is None:
            raise ValueError(f"Couldn't parse rule: {rule!r}")
        colour = m.group(1)
        child_nums = m.captures(2)
        child_colours = m.captures(3)

        children = {}

        for n, c in zip(child_nums, child_colours):
            if n == 'no':
                continue
            children[c] = int(n)

        rules_dict[colour] = children

    return rules_dict

def parents_of(rules, colour):
    parents = set()
    for parent, children in rules.items():
        if colour in children:
            parents.add(parent)
            parents |= parents_of(rules, parent)
    return parents

def children_of(rules, colour):
    children = defaultdict(lambda: 0)
    for child_colour, child_num in rules[colour].items():
        children[child_colour] += child_num
        for col, num in children_of(rules, child_colour).items():
            children[col] += child_num * num
    return children

if __name__ == '__main__':
    bag_rules = [rule.strip() for rule in sys.stdin.readlines()]
    rules = process_rules(bag_rules)
    print(len(parents_of(rules, 'shiny gold')))
    print(sum(children_of(rules, 'shiny gold').values()))