aoc-2020/7.py

# Read lines
lines = []

with open("7.input") as f:
    for line in f.readlines():
        lines.append(line.strip().replace(".", ""))

# Parse lines
bags = {}

for line in lines:
    words = line.split(" ")
    name = " ".join(words[0:2])
    bags[name] = []
    children = " ".join(words[4:]).split(",")
    if children == ["no other bags"]:
        continue
    for child in children:
        child = child.strip()
        words = child.split(" ")
        bags[name].append((int(words[0]), " ".join(words[1:3])))

# Create inverse lookup table
bag_parents = {}

for (parent, children) in bags.items():
    for (_, child) in children:
        parents = bag_parents.get(child, [])
        parents.append(parent)
        bag_parents[child] = parents

# Find all parents
flatten = lambda t: [item for sublist in t for item in sublist]

def find_parents(child):
    return flatten([[p] + find_parents(p) for p in bag_parents.get(child, [])])

print("Answer 1: {}".format(len(set(find_parents("shiny gold")))))

# Find all children
def count_children(parent):
    return sum([c * (count_children(child) + 1) for (c, child) in bags[parent]])

print("Answer 2: {}".format(count_children("shiny gold")))