16 simplify a bit

1 files changed, 11 insertions, 18 deletions

diff --git a/16.py b/16.py
index 5ea6748..d56f5bb 100644
--- a/16.py
+++ b/16.py
@@ -46,30 +46,23 @@ def recurse(open_valves, current, flow, time_left):
     cflow = sum_flow(open_valves)
     cnode = nodes[current]
     best = 0
-    for neighbour, cost in cnode.neighbours:
-        if cost >= time_left:
-            best = max(best, flow + cflow * time_left)
-        else:
-            best = max(
-                best,
-                recurse(open_valves, neighbour,
-                        flow + cflow * cost, time_left - cost)
-            )
+    def loop():
+        nonlocal best
+        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)
+    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
-        for neighbour, cost in cnode.neighbours:
-            if cost >= time_left:
-                best = max(best, flow + cflow * time_left)
-            else:
-                best = max(
-                    best,
-                    recurse(open_valves, neighbour,
-                            flow + cflow * cost, time_left - cost)
-                )
+        loop()
     return best
 
 print(recurse(frozenset(), 'AA', 0, 30))