diff options
Diffstat (limited to '2020/day7/inside_shiny.py')
-rw-r--r-- | 2020/day7/inside_shiny.py | 45 |
1 files changed, 45 insertions, 0 deletions
diff --git a/2020/day7/inside_shiny.py b/2020/day7/inside_shiny.py new file mode 100644 index 0000000..00df40b --- /dev/null +++ b/2020/day7/inside_shiny.py | |||
@@ -0,0 +1,45 @@ | |||
1 | import networkx as nx | ||
2 | import re | ||
3 | |||
4 | |||
5 | def fits_in(root): | ||
6 | """recursively count the number of bags that fit in other bags | ||
7 | |||
8 | :root: the node to count bags from | ||
9 | :returns: total | ||
10 | |||
11 | """ | ||
12 | total = 1 | ||
13 | for neighbour in bagtree[root]: | ||
14 | total += int(bagtree[root][neighbour]["weight"]) * fits_in(neighbour) | ||
15 | return total | ||
16 | |||
17 | |||
18 | bagtree = nx.DiGraph() | ||
19 | bagremover = re.compile(r" bags?\.?$") | ||
20 | countandbag = re.compile(r"(\d+) (\w+ \w+)") | ||
21 | |||
22 | with open("input", "r") as baglines: | ||
23 | for line in baglines: | ||
24 | (miniroot, child_str) = list(map(str.strip, line.split("contain"))) | ||
25 | |||
26 | miniroot = miniroot.replace(" bags", "") | ||
27 | |||
28 | children = list( | ||
29 | map( | ||
30 | lambda a: re.sub(bagremover, "", a), | ||
31 | list(map(str.strip, child_str.split(","))), | ||
32 | ) | ||
33 | ) | ||
34 | |||
35 | if "no other" in children: | ||
36 | continue | ||
37 | |||
38 | for kid in children: | ||
39 | matches = re.match(countandbag, kid) | ||
40 | bagtree.add_edge(miniroot, matches.groups()[1], weight=matches.groups()[0]) | ||
41 | |||
42 | lengths = dict(nx.all_pairs_shortest_path(bagtree)) | ||
43 | |||
44 | # we don't count the shiny gold itself | ||
45 | print(fits_in("shiny gold") - 1) | ||