16 much faster

1 files changed, 9 insertions, 7 deletions

diff --git a/16.py b/16.py
index c004091..24429a7 100644
--- a/16.py
+++ b/16.py
@@ -51,20 +51,22 @@ def sum_flow(open_valves):
         open_valves >>= 1
     return total
 
-@cache
-def recurse(open_valves=0, current=0, flow=0, time_left=30):
-    cflow = sum_flow(open_valves)
+recurse_cache = dict()
+def recurse(open_valves=1, current=0, flow=0, time_left=30):
+    key = (current, flow, time_left)
+    if key in recurse_cache: return recurse_cache[key]
     cnode = nodes[current]
-    best = flow + cflow * time_left
+    best = flow
     def check(open_valves, current, flow, time_left):
         nonlocal best
         new = recurse(open_valves, current, flow, time_left)
         if new > best: best = new
     for neighbour, cost in cnode.neighbours:
         if cost >= time_left: continue
-        check(open_valves, neighbour, flow + cflow * cost, time_left - cost)
-    if not (open_valves & (1 << current)) and time_left > 0 and cnode.flow > 0:
-        check(open_valves | 1 << current, current, flow + cflow, time_left - 1)
+        check(open_valves, neighbour, flow, time_left - cost)
+    if not (open_valves & (1 << current)) and time_left > 0:
+        check(open_valves | 1 << current, current, flow + cnode.flow * (time_left - 1), time_left - 1)
+    recurse_cache[key] = best
     return best
 
 print(recurse())