Solve 2024 till day 18

2024/d16.py

@@ -0,0 +1,96 @@
+import heapq
+from collections import defaultdict
+from lib import get_data
+from lib import Grid2D
+
+
+data = get_data(__file__)
+
+g = Grid2D(data)
+
+
+def search(start, end, dir):
+    paths = set()
+
+    def reconstruct(parents, current):
+        seen = set()
+        to_visit = [current]
+        while to_visit:
+            current = to_visit.pop()
+            if current in seen:
+                continue
+            seen.add(current)
+            paths.add(current[0])
+            for p in parents[current]:
+                to_visit.append(p)
+
+    def dist(n1, n2):
+        """cost from node to node"""
+        if n1 == 0 or n1 == n2:
+            return 0
+        p1, d1 = n1
+        p2, d2 = n2
+        if p1 != p2:
+            return 1
+        if d1 != d2:
+            return 1000
+        assert False
+
+    def h(node):
+        """heuristic function (never overestimate)"""
+        pos = node[0]
+        return abs(pos[0] - end[0]) + abs(pos[1] - end[1])
+
+    def is_goal(node):
+        return node[0] == end
+
+    def neighbors(node):
+        pos, dir = node
+        npos = pos[0] + dir[0], pos[1] + dir[1]
+        nbs = []
+        if g.contains(npos) and g[npos] != "#":
+            nbs.append((npos, dir))
+        ndir = dir[1], -dir[0]
+        nbs.append((pos, ndir))
+        ndir = -dir[1], dir[0]
+        nbs.append((pos, ndir))
+        return nbs
+
+    starts = [(start, dir)]
+    open_set = []
+    g_score = {}
+    cost = None
+
+    for start in starts:
+        heapq.heappush(open_set, (h(start), start))
+        g_score[start] = dist(0, start)
+
+    parents = defaultdict(list)
+    while open_set:
+        current_f_score, current = heapq.heappop(open_set)
+        if is_goal(current):
+            # assert current_f_score == g_score[current]
+            gs = g_score[current]
+            if cost is None or gs <= cost:
+                cost = gs
+                reconstruct(parents, current)
+            else:
+                break
+
+        for neighbor in neighbors(current):
+            tentative_g_score = g_score[current] + dist(current, neighbor)
+            if tentative_g_score <= g_score.get(neighbor, 10**12):
+                parents[neighbor].append(current)
+                g_score[neighbor] = tentative_g_score
+                f_score = g_score[neighbor] + h(neighbor)
+                heapq.heappush(open_set, (f_score, neighbor))
+    return cost, paths
+
+
+(start,) = g.find("S")
+(end,) = g.find("E")
+dir = (0, 1)
+
+cost, paths = search(start, end, dir)
+print(cost)
+print(len(paths))