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

51 lines
1.4 KiB
Python
Raw Normal View History

2020-12-07 07:15:44 +01:00
#!/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
2020-12-07 07:20:34 +01:00
def parents_of(rules, colour):
2020-12-07 07:15:44 +01:00
parents = set()
for parent, children in rules.items():
if colour in children:
parents.add(parent)
2020-12-07 07:20:34 +01:00
parents |= parents_of(rules, parent)
2020-12-07 07:15:44 +01:00
return parents
2020-12-07 07:20:34 +01:00
def children_of(rules, colour):
2020-12-07 07:15:44 +01:00
children = defaultdict(lambda: 0)
for child_colour, child_num in rules[colour].items():
children[child_colour] += child_num
2020-12-07 07:20:34 +01:00
for col, num in children_of(rules, child_colour).items():
2020-12-07 07:15:44 +01:00
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)
2020-12-07 07:20:34 +01:00
print(len(parents_of(rules, 'shiny gold')))
print(sum(children_of(rules, 'shiny gold').values()))