from __future__ import annotations import heapq from math import inf from typing import Hashable, TypeVar Node = TypeVar("Node", bound=Hashable) Graph = dict[Node, list[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.""" if start not in graph: raise KeyError(f"Unknown start node: {start!r}") distances = {node: inf for node in graph} predecessors: dict[Node, Node | None] = {node: None for node in graph} distances[start] = 0.0 queue: list[tuple[float, Node]] = [(0.0, start)] while queue: distance, node = heapq.heappop(queue) # Ignore entries made obsolete by a shorter route. if distance != distances[node]: continue for neighbor, weight in graph[node]: if weight < 0: raise ValueError("Dijkstra's algorithm requires non-negative weights") if neighbor not in graph: raise KeyError(f"Unknown neighbor node: {neighbor!r}") candidate = distance + weight if candidate < distances[neighbor]: distances[neighbor] = candidate predecessors[neighbor] = node heapq.heappush(queue, (candidate, neighbor)) return distances, predecessors def shortest_path( predecessors: dict[Node, Node | None], start: Node, target: Node ) -> list[Node]: """Reconstruct a shortest path, or return an empty list if unreachable.""" path: list[Node] = [] current: Node | None = target while current is not None: path.append(current) if current == start: return path[::-1] current = predecessors[current] return [] def main() -> None: graph: Graph[str] = { "A": [("B", 4), ("C", 1)], "B": [("D", 1)], "C": [("B", 2), ("D", 5)], "D": [], "E": [], } distances, predecessors = dijkstra(graph, "A") assert distances["D"] == 4 assert shortest_path(predecessors, "A", "D") == ["A", "C", "B", "D"] assert distances["E"] == inf assert shortest_path(predecessors, "A", "E") == [] if __name__ == "__main__": main()