2024 day 16, networkx version.

src/holt59/aoc/2024/day16.py

@@ -1,7 +1,58 @@
-from typing import Any, Iterator
+from typing import Any, Iterator, TypeAlias
+
+import networkx as nx
 
 from ..base import BaseSolver
 
+Position: TypeAlias = tuple[int, int]
+Direction: TypeAlias = tuple[int, int]
+
 
 class Solver(BaseSolver):
-    def solve(self, input: str) -> Iterator[Any]: ...
+    def solve(self, input: str) -> Iterator[Any]:
+        grid = [list(r) for r in input.splitlines()]
+        n_rows, n_cols = len(grid), len(grid[0])
+        rein = next(
+            (i, j) for i in range(n_rows) for j in range(n_cols) if grid[i][j] == "S"
+        )
+        target = next(
+            (i, j) for i in range(n_rows) for j in range(n_cols) if grid[i][j] == "E"
+        )
+        facing = (0, 1)
+
+        G = nx.DiGraph()
+
+        for i in range(1, n_rows - 1):
+            for j in range(1, n_cols - 1):
+                if grid[i][j] == "#":
+                    continue
+
+                for di, dj in ((-1, 0), (0, 1), (1, 0), (0, -1)):
+                    if grid[i + di][j + dj] != "#":
+                        G.add_edge(
+                            ((i, j), (di, dj)), ((i + di, j + dj), (di, dj)), weight=1
+                        )
+
+                    if di == 0:
+                        G.add_edge(((i, j), (di, dj)), ((i, j), (1, 0)), weight=1000)
+                        G.add_edge(((i, j), (di, dj)), ((i, j), (-1, 0)), weight=1000)
+                    else:
+                        G.add_edge(((i, j), (di, dj)), ((i, j), (0, 1)), weight=1000)
+                        G.add_edge(((i, j), (di, dj)), ((i, j), (0, -1)), weight=1000)
+
+        for di, dj in ((-1, 0), (0, 1), (1, 0), (0, -1)):
+            G.add_edge((target, (di, dj)), (target, (-1, -1)), weight=0)
+
+        yield nx.shortest_path_length(
+            G,
+            source=(rein, facing),
+            target=(target, (-1, -1)),
+            weight="weight",
+        )
+
+        nodes: set[tuple[int, int]] = set()
+        for path in nx.all_shortest_paths(
+            G, source=(rein, facing), target=(target, (-1, -1)), weight="weight"
+        ):
+            nodes.update(position for position, _ in path)
+        yield len(nodes)