https://adventofcode.com/2023/day/17

day17.py

from heapq import *

L = [[*map(int, l.strip())] for l in open(0)]
N, M = len(L), len(L[0])
MS = [(1, 0), (-1, 0), (0, 1), (0, -1)]

x, y = 0, 0
E = (N-1, M-1)


def dijkstra(P2):
    Q, V = [(0, 0, 0, -1, 0)], set()
    while Q:
        c, x, y, d, l = heappop(Q)
        if (x, y, d, l) in V: continue
        V.add((x, y, d, l))
        if (x, y)==E: return c

        for nd in range(4):
            if nd==d and l<=0 or (d, nd) in [(0, 1), (1, 0), (2, 3), (3, 2)]:
                continue
            if P2 and nd!=d and l>6: continue
            dx, dy = MS[nd]
            nx, ny = x+dx, y+dy
            if nx<0 or nx>=N or ny<0 or ny>=M: continue
            heappush(Q, (c+L[nx][ny], nx, ny, nd, l-1 if nd==d else
                         (9 if P2 else 2)))
    return 1e9


print(dijkstra(False))
print(dijkstra(True))