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)