no turning back

day17/sol.py

@@ -11,12 +11,22 @@ def solve(map,part=1):
     goal = (len(map)-1,len(map[-1])-1)
     def is_goal(x):
         return x[:2] == goal
+    def heuristic(state):
+        # manhattan distance. this is correct because the cost
+        # at every point is >= 1 and we can only move one square
+        # at a time NSEW.
+        if part == 2 and state[3] < 4:
+            y,x,d,count = state
+            # account for forced moves away from the goal
+            if d[0] < 0 or d[1] < 0:
+                return (4-count)*2 + abs(y-goal[0]) + abs(x-goal[1])
+        return abs(state[0]-goal[0]) + abs(state[1]-goal[1])
     def neighbors(state):
         y,x,dir,count = state
         next = []
         def go(dx,dy):
             if (-dx,-dy) == dir:
-                # can't turn around
+                # no turning back
                 return
             newx = x + dx
             newy = y + dy
@@ -40,6 +50,7 @@ def solve(map,part=1):
         go(0,+1)
         return next
     
+    # turns out it's actually faster to not use the heuristic, so don't
     cost, node, path = astar.search(start, is_goal, neighbors)
     print(cost)
     print([(x,y) for (y,x,_,_) in path])