48 lines
1.3 KiB
Python
48 lines
1.3 KiB
Python
#!/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")
|