51 lines
1.4 KiB
Python
51 lines
1.4 KiB
Python
|
#!/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 can_contain(rules, colour):
|
||
|
|
parents = set()
|
||
|
|
for parent, children in rules.items():
|
||
|
|
if colour in children:
|
||
|
|
parents.add(parent)
|
||
|
|
parents |= can_contain(rules, parent)
|
||
|
|
return parents
|
||
|
|
|
||
|
|
def colour_contains(rules, colour):
|
||
|
|
children = defaultdict(lambda: 0)
|
||
|
|
for child_colour, child_num in rules[colour].items():
|
||
|
|
children[child_colour] += child_num
|
||
|
|
for col, num in colour_contains(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(can_contain(rules, 'shiny gold')))
|
||
|
|
print(sum(colour_contains(rules, 'shiny gold').values()))
|