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)

day-16/day-16.py

@@ -73,38 +73,23 @@ def calc_result(packets):
     return res
 
 
-def part_1(input):
-    result = 0
+def solve(input):
+    result_p1 = 0
     line = input[0].rstrip().strip('0')
-    msg = bin(int(line, 16))
-    msg = msg[2:]
-    while len(msg) % 4:
+    msg = bin(int(line, 16))[2:]
+    while len(msg) < len(line) * 4:
         msg = '0' + msg
     packets = []
     next = 0
     while -1 != next:
-        next, result = get_next_packet(msg, next, packets, result)
-    print("Part 1 result:", result)
-
-
-def part_2(input):
-    result = 0
-    line = input[0].rstrip().strip('0')
-    msg = bin(int(line, 16))
-    msg = msg[2:]
-    while len(msg) % 4:
-        msg = '0' + msg
-    packets = []
-    next = 0
-    while -1 != next:
-        next, _ = get_next_packet(msg, next, packets)
-    result = calc_result(packets)
-    print("Part 2 result:", result)
+        next, result_p1 = get_next_packet(msg, next, packets, result_p1)
+    result_p2 = calc_result(packets)
+    print("Part 1 result:", result_p1)
+    print("Part 2 result:", result_p2)
 
 
 input = list()
 p = Path(__file__).with_name('input.txt')
 with open(p) as f:
     input = f.readlines()
-    part_1(input)
-    part_2(input)
+    solve(input)