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)
|