Add generic simple dijkstra method.

src/holt59/aoc/tools/graphs.py

@@ -0,0 +1,95 @@
+import heapq
+from typing import Callable, Iterable, TypeVar
+
+_Node = TypeVar("_Node")
+
+
+def make_neighbors_grid_fn(
+    rows: int | Iterable[int],
+    cols: int | Iterable[int],
+    excluded: Iterable[tuple[int, int]] = set(),
+    diagonals: bool = False,
+):
+    """
+    Create a neighbors function suitable for graph function for a simple grid.
+
+    Args:
+        rows: Rows of the grid. If an int is specified, the rows are assumed to be
+            numbered from 0 to rows - 1, otherwise the iterable should contain the list
+            of valid rows.
+        cols: Columns of the grid. If an int is specified, the columns are assumed to be
+            numbered from 0 to cols - 1, otherwise the iterable should contain the list
+            of valid columns.
+        excluded: Cells of the grid that cannot be used as valid nodes for the graph.
+        diagonals: If True, neighbors will include diagonal cells, otherwise, only
+            horizontal and vertical neighbors will be included.
+
+    """
+    ds = ((-1, 0), (0, 1), (1, 0), (0, -1))
+    if diagonals:
+        ds = ds + ((-1, -1), (-1, 1), (1, -1), (1, 1))
+
+    if isinstance(rows, int):
+        rows = range(rows)
+    elif not isinstance(rows, range):
+        rows = set(rows)
+
+    if isinstance(cols, int):
+        cols = range(cols)
+    elif not isinstance(cols, range):
+        cols = set(cols)
+
+    excluded = set(excluded)
+
+    def _fn(node: tuple[int, int]):
+        return (
+            ((row_n, col_n), 1)
+            for dr, dc in ds
+            if (row_n := node[0] + dr) in rows
+            and (col_n := node[1] + dc) in cols
+            and (row_n, col_n) not in excluded
+        )
+
+    return _fn
+
+
+def dijkstra(
+    start: _Node,
+    target: _Node,
+    neighbors: Callable[[_Node], Iterable[tuple[_Node, float]]],
+) -> tuple[tuple[_Node, ...], float] | None:
+    """
+    Solve shortest-path problem using simple Dijkstra algorithm from start to target,
+    using the given neighbors function.
+
+    Args:
+        start: Starting node of the path.
+        target: Target node for the path.
+        neighbors: Function that should return, for a given node, the list of
+            its neighbors with the cost to go from the node to the neighbor.
+
+    Returns:
+        One of the shortest-path from start to target with its associated cost, if one
+        is found, otherwise None.
+    """
+    queue: list[tuple[float, _Node, tuple[_Node, ...]]] = [(0, start, (start,))]
+    preds: dict[_Node, tuple[tuple[_Node, ...], float]] = {}
+
+    while queue:
+        dis, node, path = heapq.heappop(queue)
+
+        if node in preds:
+            continue
+
+        preds[node] = (path, dis)
+
+        if node == target:
+            break
+
+        for neighbor, cost in neighbors(node):
+            if neighbor in preds:
+                continue
+
+            heapq.heappush(queue, (dis + cost, neighbor, path + (neighbor,)))
+
+    return preds.get(target, None)