aoc

Advent of Code
git clone https://git.pngu.org/aoc
Log | Files | Refs | README

7_2020.py (1944B)


import fileinput
import re
from collections import deque, defaultdict

# For part1, we need to keep track of which bags have other bags as containers
# For part2, we need to keep track of which bags (and how many!) each bag contains
# We can think of both of these as graphs where the vertices are the bags and the edges
#  are which bags might contain us (part1) and which bags and how many of each we contain (part2)

PARENTS = defaultdict(list)  # PARENTS[x] are the bags that contain x
CONTENTS = defaultdict(
    list
)  # CONTENTS[x] are which bags x contains and how many of each

target = "shinygoldbag"

lines = list(fileinput.input())
lines.append("")
for line in lines:
    line = line.strip()
    if line:
        words = line.split()
        container = words[0] + words[1] + words[2]
        container = container[:-1]  # remove trailing 's' in 'bags'
        if words[-3] == "no":  # doesn't contain any other bags
            continue
        else:
            idx = 4
            while idx < len(words):
                bag = words[idx] + words[idx + 1] + words[idx + 2] + words[idx + 3]
                if bag.endswith("."):
                    bag = bag[:-1]
                if bag.endswith(","):
                    bag = bag[:-1]
                if bag.endswith("s"):
                    bag = bag[:-1]
                n = int(bag[0])
                assert bag[1] not in "0123456789"
                while any([bag.startswith(d) for d in "0123456789"]):
                    bag = bag[1:]

                PARENTS[bag].append(container)
                CONTENTS[container].append((n, bag))
                idx += 4

SEEN = set()
Q = deque([target])
while Q:
    x = Q.popleft()
    if x in SEEN:
        continue
    SEEN.add(x)
    for y in PARENTS[x]:
        Q.append(y)
print(len(SEEN) - 1)


def size(bag):
    ans = 1
    for n, y in CONTENTS[bag]:
        ans += n * size(y)
    return ans


print(size(target) - 1)