49 lines
1.1 KiB
Python
49 lines
1.1 KiB
Python
import fileinput
|
|
import re
|
|
from collections import defaultdict
|
|
|
|
network = {}
|
|
reverse = defaultdict(list)
|
|
|
|
for line in fileinput.input():
|
|
parent, children = line.strip().split(" contain ")
|
|
parent = parent[:-5]
|
|
children = children.rstrip(".").split(", ")
|
|
children = [
|
|
child.rstrip("s")[:-4].split(" ", 1)
|
|
for child in children
|
|
if "no other" not in child
|
|
]
|
|
children = [(int(n), c) for n, c in children]
|
|
|
|
network[parent] = children
|
|
for n, child in children:
|
|
reverse[child].append((n, parent))
|
|
|
|
reverse = dict(reverse)
|
|
|
|
shiny_gold_bags = set()
|
|
stack = ["shiny gold"]
|
|
while stack:
|
|
color = stack.pop()
|
|
if color not in reverse:
|
|
continue
|
|
for _, parent in reverse[color]:
|
|
stack.append(parent)
|
|
shiny_gold_bags.add(parent)
|
|
|
|
print("part 1:", len(shiny_gold_bags))
|
|
|
|
part2 = 0
|
|
stack = [(1, "shiny gold")]
|
|
while stack:
|
|
n, color = stack.pop()
|
|
part2 += n
|
|
|
|
children = network[color]
|
|
for m, child in children:
|
|
stack.append((n * m, child))
|
|
part2 -= 1
|
|
|
|
print("part 2:", part2)
|