54 lines
1.3 KiB
Python
54 lines
1.3 KiB
Python
import re
|
|
from pprint import pprint
|
|
|
|
|
|
def load_rules(filename):
|
|
re_rule = re.compile(r"(?P<color>.+?) bags contain (?P<contained>.+)")
|
|
re_contents = re.compile(r"(?P<num_bags>\d+)\s+(?P<color>.+?)\s+bags?")
|
|
|
|
rules = {}
|
|
with open(filename, "r", encoding="utf-8") as f:
|
|
for line in f:
|
|
m = re_rule.match(line.rstrip())
|
|
assert m, line
|
|
color = m.group("color")
|
|
contained = m.group("contained")
|
|
rules[color] = {}
|
|
if contained == "no other bags.":
|
|
continue
|
|
|
|
m = re_contents.finditer(contained)
|
|
rules[color] = { x.group("color"): int(x.group("num_bags")) for x in m }
|
|
|
|
return rules
|
|
|
|
|
|
def what_can_hold(query_color, rules):
|
|
for outer_color, outer_contents in rules.items():
|
|
if not outer_contents:
|
|
continue
|
|
|
|
for inner_color, num_bags in outer_contents.items():
|
|
if inner_color == query_color:
|
|
yield outer_color
|
|
yield from what_can_hold(outer_color, rules)
|
|
|
|
|
|
def bag_contains(query_color, rules):
|
|
rule = rules[query_color]
|
|
if not rule:
|
|
return 0
|
|
|
|
x = 0
|
|
for color, multiplier in rule.items():
|
|
x += multiplier * (bag_contains(color, rules) + 1)
|
|
|
|
return x
|
|
|
|
rules = load_rules("data.txt")
|
|
|
|
wch = set(what_can_hold("shiny gold", rules))
|
|
print(f"Shiny gold can be held by {len(wch)} bag colors.")
|
|
|
|
bc = bag_contains("shiny gold", rules)
|
|
print(f"A single shiny gold bag requires {bc} individual bags.")
|