summaryrefslogtreecommitdiffstats
path: root/2020/day7/inside_shiny.py
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)