from __future__ import annotations import heapq from collections.abc import Hashable, Mapping, Sequence from math import inf from typing import TypeVar Node = TypeVar("Node", bound=Hashable) Graph = Mapping[Node, Sequence[tuple[Node, float]]] def dijkstra(graph: Graph[Node], start: Node) -> tuple[dict[Node, float], dict[Node, Node | None]]: """Return shortest distances and predecessors from start.""" nodes = set(graph) for edges in graph.values(): for neighbor, weight in edges: if weight < 0: raise ValueError("Dijkstra's algorithm requires non-negative weights") nodes.add(neighbor) if start not in nodes: raise KeyError(f"Unknown start node: {start!r}") distances = {node: inf for node in nodes} predecessors: dict[Node, Node | None] = {node: None for node in nodes} distances[start] = 0.0 # A counter avoids comparing nodes when equal distances are queued. counter = 0 queue: list[tuple[float, int, Node]] = [(0.0, counter, start)] while queue: distance, _, node = heapq.heappop(queue) if distance != distances[node]: continue # Ignore stale queue entries. for neighbor, weight in graph.get(node, ()): candidate = distance + weight if candidate < distances[neighbor]: distances[neighbor] = candidate predecessors[neighbor] = node counter += 1 heapq.heappush(queue, (candidate, counter, neighbor)) return distances, predecessors def shortest_path( graph: Graph[Node], start: Node, target: Node ) -> tuple[float, list[Node]]: """Return the shortest distance and path, or (inf, []) if unreachable.""" distances, predecessors = dijkstra(graph, start) if target not in distances: raise KeyError(f"Unknown target node: {target!r}") if distances[target] == inf: return inf, [] path: list[Node] = [] current: Node | None = target while current is not None: path.append(current) current = predecessors[current] path.reverse() return distances[target], path def main() -> None: graph: Graph[str] = { "A": [("B", 4), ("C", 1)], "B": [("D", 1)], "C": [("B", 2), ("D", 5)], "D": [], "E": [], } distance, path = shortest_path(graph, "A", "D") assert distance == 4 assert path == ["A", "C", "B", "D"] distance, path = shortest_path(graph, "A", "E") assert distance == inf assert path == [] if __name__ == "__main__": main()