summaryrefslogtreecommitdiff
path: root/7_2020.py
blob: ede07f3e845ea460a01e8042cddd377cc913854d (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
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)