16 shorter

1 files changed, 11 insertions, 17 deletions

diff --git a/16.py b/16.py
index d56f5bb..acddc1f 100644
--- a/16.py
+++ b/16.py
@@ -25,15 +25,11 @@ for valve, (flow, neighbours) in inp.items():
     for n in neighbours:
         prev = valve
         for cost in count(1):
-            if n == 'AA' or inp[n][0] != 0:
-                break
+            if n == 'AA' or inp[n][0] != 0: break
             l, r = inp[n][1]
-            if l == prev:
-                prev = n
-                n = r
-            else:
-                prev = n
-                n = l
+            nnext = r if l == prev else l
+            prev = n
+            n = nnext
         actual_neighbours.append((n, cost))
     nodes[valve] = Node(flow, actual_neighbours)
 
@@ -42,27 +38,25 @@ def sum_flow(open_valves):
     return sum(nodes[n].flow for n in open_valves)
 
 @cache
-def recurse(open_valves, current, flow, time_left):
+def recurse(open_valves=frozenset(), current='AA', flow=0, time_left=30):
     cflow = sum_flow(open_valves)
     cnode = nodes[current]
     best = 0
     def loop():
         nonlocal best
+        best = max(best, flow + cflow * time_left)
         for neighbour, cost in cnode.neighbours:
-            if cost >= time_left:
-                best = max(best, flow + cflow * time_left)
-            else:
-                new = recurse(open_valves, neighbour,
-                              flow + cflow * cost, time_left - cost)
-                best = max(best, new)
+            if cost >= time_left: continue
+            new = recurse(open_valves, neighbour,
+                          flow + cflow * cost, time_left - cost)
+            best = max(best, new)
     loop()
     if current not in open_valves and time_left > 0 and cnode.flow > 0:
         flow += cflow
         open_valves |= {current}
         cflow = sum_flow(open_valves)
-        best = max(best, flow)
         time_left -= 1
         loop()
     return best
 
-print(recurse(frozenset(), 'AA', 0, 30))
+print(recurse())