adventofcode2023/day17/sol.py

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)