AoC20/d07.py

#!/usr/bin/env python3

from sys import argv

i = [x.strip() for x in open(argv[1])]

b = {}
r = {}
for l in i:
    c, l = l.split(" bags contain ")
    if c not in b:
        b[c] = {}
    if l.startswith("no other bags"):
        continue
    l = l.split(", ")
    for i in l:
        i = i.split(" ")
        n = int(i[0])
        cc = " ".join(i[1:3])
        b[c][cc] = n
        if cc not in r:
            r[cc] = set()
        r[cc].add(c)

o = set()
def add(x):
    global o
    o.add(x)
    if x in r:
        for c in r[x]:
            if c not in o:
                add(c)
add("shiny gold")
print(len(o)-1)

ch = {}
def add(x):
    o = 1
    for c in b[x]:
        if c in ch:
            o += ch[c]*b[x][c]
        else:
            o += add(c)*b[x][c]
    ch[x] = o
    return o
o = add("shiny gold")
print(o-1)