Optimize day 15 solution by using priority queue

day-15/day-15.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 from pathlib import Path
+from heapq import heappush, heappop
 
 
 def part_1(input):
@@ -11,16 +12,13 @@ def part_1(input):
     start = (0, 0)
     end = (x_size - 1, y_size - 1)
     cost = {start: 0}
-    queue = set([start])
+    queue = [(0, start)]
     done = set()
     neighbors = set([(-1, 0), (0, -1), (0, 1), (1, 0)])
     while len(queue):
-        cur_cost = None
-        for n in queue:
-            if cur_cost == None or cost[n] < cur_cost:
-                cur_cost = cost[n]
-                current = n
-        queue.remove(current)
+        cur_cost, current = heappop(queue)
+        if current in done:
+            continue
         (xc, yc) = current
         next = set()
         for (xn, yn) in neighbors:
@@ -31,7 +29,7 @@ def part_1(input):
             next_cost = cur_cost + nodes[n]
             if not n in cost or next_cost < cost[n]:
                 cost[n] = next_cost
-            queue.add(n)
+                heappush(queue, (next_cost, n))
         done.add(current)
     result = cost[end]
     print("Part 1 result:", result)
@@ -47,16 +45,13 @@ def part_2(input):
     start = (0, 0)
     end = ((x_size * repeat) - 1, (y_size * repeat) - 1)
     cost = {start: 0}
-    queue = set([start])
+    queue = [(0, start)]
     done = set()
     neighbors = set([(-1, 0), (0, -1), (0, 1), (1, 0)])
     while len(queue):
-        cur_cost = None
-        for n in queue:
-            if cur_cost == None or cost[n] < cur_cost:
-                cur_cost = cost[n]
-                current = n
-        queue.remove(current)
+        cur_cost, current = heappop(queue)
+        if current in done:
+            continue
         (xc, yc) = current
         next = set()
         for (xn, yn) in neighbors:
@@ -72,7 +67,7 @@ def part_2(input):
                 next_cost -= 9
             if not (x, y) in cost or (cur_cost + next_cost) < cost[(x, y)]:
                 cost[(x, y)] = cur_cost + next_cost
-            queue.add((x, y))
+                heappush(queue, (cur_cost + next_cost, (x, y)))
         done.add(current)
     result = cost[end]
     print("Part 2 result:", result)