day 24 python part 2

day24/sol.py

@@ -54,9 +54,6 @@ def show(t):
         print(''.join(s))
 
 def blocked(x,y,t):
-    H = len(data)-2
-    W = len(data[y])-2
-
     if not 0 <= y < len(data):
         return True
     if not 0 <= x < len(data[y]):
@@ -64,6 +61,9 @@ def blocked(x,y,t):
     if data[y][x] == '#':
         return True
 
+    H = len(data)-2
+    W = len(data[y])-2
+
     u = 1 + ((x - t)-1)%W
     if (u,y) in blizzards['>']:
         return True
@@ -82,23 +82,35 @@ def blocked(x,y,t):
 
     return False
 
-start = (data[0].index('.'), 0, 0)
+start = (data[0].index('.'), 0)
 goal = (data[-1].index('.'), len(data)-1)
+leg_distance = abs(start[0] - goal[0]) + abs(start[1] - goal[1])
 
 def is_goal(node):
-    x, y, t = node
-    return (x,y) == goal
+    x, y, t, legs = node
+    return (x,y) == goal and legs == 0
 
 def heuristic(node):
-    x, y, t = node
-    return abs(goal[0] - x) + abs(goal[1] - y)
+    x, y, t, legs = node
+    if legs == 0:
+        return 0
+    if legs % 2 == 1:
+        g = goal
+    else:
+        g = start
+    return (legs-1)*leg_distance + abs(g[0] - x) + abs(g[1] - y)
 
 def neighbors(node):
-    x, y, t = node
+    x, y, t, legs = node
     n = []
+    if legs % 2 == 1:
+        g = goal
+    else:
+        g = start
     def check(dx,dy):
         if not blocked(x+dx,y+dy,t+1):
-            n.append((1, (x+dx, y+dy, t+1)))
+            l = legs - ((x+dx,y+dy) == g)
+            n.append((1, (x+dx, y+dy, t+1, l)))
     check(+1,0)
     check(0,+1)
     check(-1,0)
@@ -110,4 +122,8 @@ show(0)
 show(1)
 #show((len(data)-2)*(len(data[0])-2))
 
-print(astar.search(start, is_goal, neighbors, heuristic))
+start_node = start + (0, 1)
+print(astar.search(start_node, is_goal, neighbors, heuristic))
+
+start_node = start + (0, 3)
+print(astar.search(start_node, is_goal, neighbors, heuristic))