70 lines
2.0 KiB
Python
70 lines
2.0 KiB
Python
import astar
|
|
|
|
def read_input(f):
|
|
map = []
|
|
for line in f:
|
|
map.append([int(x) for x in line.strip()])
|
|
return map
|
|
|
|
def solve(map,part=1):
|
|
start = (0,0,(0,0),0)
|
|
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:
|
|
# no turning back
|
|
return
|
|
newx = x + dx
|
|
newy = y + dy
|
|
if (dx,dy) == dir:
|
|
newcount = count+1
|
|
else:
|
|
newcount = 1
|
|
if part == 1:
|
|
if newcount > 3:
|
|
return
|
|
if part == 2:
|
|
if newcount > 10:
|
|
return
|
|
if newcount <= count < 4:
|
|
return
|
|
if (0 <= newy < len(map)) and (0 <= newx < len(map[newy])):
|
|
next.append((map[newy][newx], (newy, newx, (dx,dy), newcount)))
|
|
go(-1,0)
|
|
go(+1,0)
|
|
go(0,-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])
|
|
print(''.join(dir2s(d) for (_,_,d,_) in path))
|
|
|
|
def dir2s(d):
|
|
if d == (0,1): return "v"
|
|
if d == (0,-1): return "^"
|
|
if d == (1,0): return ">"
|
|
if d == (-1,0): return "<"
|
|
return "?"
|
|
|
|
import sys
|
|
map = read_input(sys.stdin)
|
|
solve(map, part=1)
|
|
solve(map, part=2)
|