From 5ae8933fca338ade5cffdc69d6eb4759a196c098 Mon Sep 17 00:00:00 2001 From: Marijn Doeve Date: Wed, 15 Dec 2021 13:52:02 +0100 Subject: [PATCH] Optimize --- 15/15.py | 21 +++++++++------------ 1 file changed, 9 insertions(+), 12 deletions(-) diff --git a/15/15.py b/15/15.py index d5b1fe4..585fcb2 100644 --- a/15/15.py +++ b/15/15.py @@ -1,10 +1,10 @@ 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: @@ -19,22 +19,18 @@ def neighbors(pos: tuple[int, int], shape): def dijkstra(grid): dist = np.full_like(grid, np.iinfo(np.int64).max) - - q = [] - - heappush(q, (0, (0, 0))) + q = [(0, (0, 0))] dist[0, 0] = 0 while q: - dist_u, u = heappop(q) - - if dist[u] < dist_u: - continue + _, u = heappop(q) 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)) + new = dist[u] + grid[neighbour] + + if dist[neighbour] > new: + dist[neighbour] = new + heappush(q, (new, neighbour)) return dist[(grid.shape[0] - 1, grid.shape[1] - 1)] @@ -59,5 +55,6 @@ if __name__ == "__main__": 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)))