15

15/15.py

@@ -0,0 +1,63 @@
+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)))