blob: 00df40b95730ac419d2b2e1ef4b860d6b1e1a64e (
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
|
import networkx as nx
import re
def fits_in(root):
"""recursively count the number of bags that fit in other bags
:root: the node to count bags from
:returns: total
"""
total = 1
for neighbour in bagtree[root]:
total += int(bagtree[root][neighbour]["weight"]) * fits_in(neighbour)
return total
bagtree = nx.DiGraph()
bagremover = re.compile(r" bags?\.?$")
countandbag = re.compile(r"(\d+) (\w+ \w+)")
with open("input", "r") as baglines:
for line in baglines:
(miniroot, child_str) = list(map(str.strip, line.split("contain")))
miniroot = miniroot.replace(" bags", "")
children = list(
map(
lambda a: re.sub(bagremover, "", a),
list(map(str.strip, child_str.split(","))),
)
)
if "no other" in children:
continue
for kid in children:
matches = re.match(countandbag, kid)
bagtree.add_edge(miniroot, matches.groups()[1], weight=matches.groups()[0])
lengths = dict(nx.all_pairs_shortest_path(bagtree))
# we don't count the shiny gold itself
print(fits_in("shiny gold") - 1)
|