jokes

y2021/p15.py

+import heapq
+
+
+def shortest_path(graph, X, Y):
+    risk = {0: 0}
+    seen = set()
+    heap = [(0, '', 0)]  # string disambiguates cost ties for imaginary numbers
+
+    end = complex(X - 1, Y - 1)
+    tip = 0
+    while end not in seen:
+        _, _, tip = heapq.heappop(heap)
+        seen.add(tip)
+        for y in (tip + dx for dx in [1, -1, 1j, -1j]):
+            if y in graph:
+                if y not in risk or risk[y] > graph[y] + risk[tip]:
+                    risk[y] = graph[y] + risk[tip]
+                    heapq.heappush(heap, (risk[y], str(y), y))
+
+    return risk[end]
+
+
+graph = {}
+for y, line in enumerate(text.splitlines()):
+    for x, char in enumerate(line):
+        graph[complex(x, y)] = int(char)
+X, Y = x + 1, y + 1
+
+ans1 = shortest_path(graph, X, Y)
+
+graph = {
+    pos + complex(X * dx, Y * dy): sum(divmod(graph[pos] + dx + dy, 10))
+    for pos, value in graph.items()
+    for dx in range(5)
+    for dy in range(5)
+}
+ans2 = shortest_path(graph, 5 * X, 5 * Y)