Solve day 16 part 2 and simplify part 1

2022/16_tunnels.py

@@ -1,39 +1,30 @@
 #!/usr/bin/env python3
 import sys
+from functools import lru_cache
 
-def parse(inp):
-    rates = {}
-    edges = {}
-    for line in inp:
-        _, node, _, _, flow, _, _, _, _, nbs = line.split(' ', 9)
-        rates[node] = int(flow[5:-1])
-        edges[node] = nbs.rstrip().split(', ')
-    return rates, edges
+def parse(line):
+    _, node, _, _, flow, _, _, _, _, nbs = line.split(' ', 9)
+    return node, int(flow[5:-1]), nbs.rstrip().split(', ')
 
-def shortest_paths(edges):
-    inf = 10000
-    dist = {node: {nb: 1 for nb in nbs} for node, nbs in edges.items()}
-    dist['start'] = {'AA': 0}
-    for node in dist:
-        dist[node][node] = 0
-    for a, adist in dist.items():
-        for b, bdist in dist.items():
-            for c in dist:
-                alt = bdist.get(a, inf) + adist.get(c, inf)
-                if bdist.get(c, inf) > alt:
-                    bdist[c] = alt
-    return dist
+def open_valves(graph, time, helpers):
+    @lru_cache(maxsize=None)
+    def travel(opened, pos, t, helpers):
+        if t == 0:
+            return travel(opened, 'AA', time, helpers - 1) if helpers else 0
 
-def open_valves(rates, dist, time):
-    def travel(released, time, pos, closed):
-        yield released
-        for dst in closed:
-            t = time - dist[pos][dst] - 1
-            if t >= 0:
-                yield from travel(released + rates[dst] * t, t, dst, closed - {dst})
+        i, rate, nbs = graph[pos]
+        mask = 1 << i
+        best = 0
+        if rate and opened & mask == 0:
+            rest = travel(opened | mask, pos, t - 1, helpers)
+            best = max(best, rate * (t - 1) + rest)
+        for nb in nbs:
+            best = max(best, travel(opened, nb, t - 1, helpers))
+        return best
 
-    valves = {node for node, rate in rates.items() if rate}
-    return max(travel(0, time, 'start', valves))
+    return travel(0, 'AA', time, helpers)
 
-rates, edges = parse(sys.stdin)
-print(open_valves(rates, shortest_paths(edges), 30))
+graph = {valve: (i, rate, nbs)
+         for i, (valve, rate, nbs) in enumerate(map(parse, sys.stdin))}
+print(open_valves(graph, 30, 0))
+print(open_valves(graph, 26, 1))