64 lines
1.4 KiB
Python
64 lines
1.4 KiB
Python
import numpy as np
|
|
from heapq import heappush, heappop
|
|
import sys
|
|
|
|
|
|
def neighbors(pos: tuple[int, int], shape):
|
|
n = []
|
|
if pos[0] > 0:
|
|
n.append((pos[0] - 1, pos[1]))
|
|
if pos[1] > 0:
|
|
n.append((pos[0], pos[1] - 1))
|
|
if pos[0] != shape[0] - 1:
|
|
n.append((pos[0] + 1, pos[1]))
|
|
if pos[1] != shape[1] - 1:
|
|
n.append((pos[0], pos[1] + 1))
|
|
|
|
return n
|
|
|
|
|
|
def dijkstra(grid):
|
|
dist = np.full_like(grid, np.iinfo(np.int64).max)
|
|
|
|
q = []
|
|
|
|
heappush(q, (0, (0, 0)))
|
|
dist[0, 0] = 0
|
|
|
|
while q:
|
|
dist_u, u = heappop(q)
|
|
|
|
if dist[u] < dist_u:
|
|
continue
|
|
|
|
for neighbour in neighbors(u, grid.shape):
|
|
if dist[neighbour] > dist[u] + grid[neighbour]:
|
|
dist[neighbour] = dist[u] + grid[neighbour]
|
|
heappush(q, (dist[neighbour], neighbour))
|
|
|
|
return dist[(grid.shape[0] - 1, grid.shape[1] - 1)]
|
|
|
|
|
|
def multiply_grid(grid: np.ndarray, n=5) -> np.ndarray:
|
|
|
|
for axis in range(2):
|
|
to_copy = grid.copy()
|
|
|
|
for i in range(1, n):
|
|
new = to_copy + i
|
|
grid = np.concatenate([grid, new], axis=axis)
|
|
|
|
grid[grid > 9] -= 9
|
|
|
|
return grid
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
grid = np.array(
|
|
[[int(x) for x in list(line.strip())] for line in sys.stdin.readlines()]
|
|
)
|
|
print("Part one:", dijkstra(grid))
|
|
print("Part two:", dijkstra(multiply_grid(grid)))
|