Advent of Code 2020 Day07 Python

aoc202007.py

from collections import defaultdict
import re

def bag(words):
    return ' '.join([words[0], words[1]])

def parse(data):
    head, tail = data.split(' contain ')
    tail = tail.replace('no other bags.', '0 other bags')
    contents = [
        re.findall(r'\s*(\w+)\s+(\w+)\s+(\w+)', ss)
        for ss in tail.split(',')
    ]
    return (bag(head.split()), [(int(w[0][0]), bag(w[0][1:])) for w in contents])

def invert(data):
    contained_in = defaultdict(set)
    for spec in data:
        for (_, name) in spec[1]:
            contained_in[name].add(spec[0])
    return contained_in

def vert(data):
    contains = defaultdict(dict)
    for spec in data:
        for (mag, name) in spec[1]:
            contains[spec[0]][name] = mag     
    return contains

def part1(inverted):
    found = set()
    queue = list(inverted['shiny gold'])
    for item in queue:
        found.add(item)
        if item in inverted:
            queue.extend(list(inverted[item]))
    return len(found)

def part2(verted):
    known = {'other bags': 0}
    while 'shiny gold' not in known:
        for k, v in verted.items():
            if set(v.keys()) <= set(known.keys()):
                known[k] = sum(vv * (1 + known[kk]) for kk, vv in v.items())
    return known['shiny gold']

with open('data/2020/day07.txt') as f:
    data = [parse(l) for l in f.read().splitlines()]

p1 = part1(invert(data))
p2 = part2(vert(data))

assert p1 == 185
assert p2 == 89084

print(p1, p2)