Add defaultdict thing back in

2023/day17/solution.py

@@ -1,4 +1,5 @@
 from heapq import heappop, heappush
+from collections import defaultdict
 
 with open("input") as f:
     inp = f.read().strip().split("\n")
@@ -10,51 +11,56 @@
 
 # state = (heat, id, position, direction_moved, moved_blocks)
 q = [(grid[p], id(p), p, p, 1) for p in [1, 1j]]
-seen = set()
+min_heat = defaultdict(lambda: float("inf"))
 
 while q:
     heat, _, p, dp, blocks = heappop(q)
-    if (p, dp, blocks) in seen:
-        continue
-    seen.add((p, dp, blocks))
     if p == corner:
         print(heat)
         break
-    dirs = [dp*1j, -dp*1j]
+    to_push = []
     # Turn.
-    for new_dp in dirs:
+    for new_dp in [dp*1j, -dp*1j]:
         new_p = p + new_dp
         if new_p in grid:
-            heappush(q, (heat + grid[new_p], id(new_p), new_p, new_dp, 1))
+            to_push.append((heat + grid[new_p], new_p, new_dp, 1))
     # If moved less than 3 blocks we can continue in the same direction.
     if blocks < 3:
         new_p = p + dp
         if new_p in grid:
-            heappush(q, (heat + grid[new_p], id(new_p), new_p, dp, blocks + 1))
+            to_push.append((heat + grid[new_p], new_p, dp, blocks + 1))
+    # Add all states to queue if heat is lower.
+    for (heat, p, dp, blocks) in to_push:
+        if heat < min_heat[(p, dp, blocks)]:
+            heappush(q, (heat, id(p), p, dp, blocks))
+            min_heat[(p, dp, blocks)] = heat
+
 
 # Part 2
 
 # state = (heat, id, position, direction_moved, moved_blocks)
 q = [(grid[p], id(p), p, p, 1) for p in [1, 1j]]
-seen = set()
+min_heat = defaultdict(lambda: float("inf"))
 
 while q:
     heat, _, p, dp, blocks = heappop(q)
-    if (p, dp, blocks) in seen:
-        continue
-    seen.add((p, dp, blocks))
     if (p == corner) & (blocks >= 4):
         print(heat)
         break
+    to_push = []
     # If moved less than 10 blocks we can continue in the same direction.
     if blocks < 10:
         new_p = p + dp
         if new_p in grid:
-            heappush(q, (heat + grid[new_p], id(new_p), new_p, dp, blocks + 1))
+            to_push.append((heat + grid[new_p], new_p, dp, blocks + 1))
     # If moved at least 4 blocks, we can turn.
     if blocks >= 4:
-        dirs = [dp*1j, -dp*1j]
-        for new_dp in dirs:
+        for new_dp in [dp*1j, -dp*1j]:
             new_p = p + new_dp
             if new_p in grid:
-                heappush(q, (heat + grid[new_p], id(new_p), new_p, new_dp, 1))
+                to_push.append((heat + grid[new_p], new_p, new_dp, 1))
+    # Add all states to queue if heat is lower.
+    for (heat, p, dp, blocks) in to_push:
+        if heat < min_heat[(p, dp, blocks)]:
+            heappush(q, (heat, id(p), p, dp, blocks))
+            min_heat[(p, dp, blocks)] = heat