advent-of-code-2k20/day7/day7.py

48 lines
1.3 KiB
Python
Raw Permalink Normal View History

2020-12-07 16:25:02 +00:00
#!/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")