35 lines
921 B
Python
35 lines
921 B
Python
"""Dijkstra shortest path on weighted directed graph.
|
|
|
|
Public API:
|
|
dijkstra(adj, src) -> dict[node, distance]
|
|
adj: dict node -> list of (neighbor, weight). Non-negative weights required.
|
|
"""
|
|
import heapq
|
|
|
|
|
|
def dijkstra(adj, src):
|
|
if src not in adj:
|
|
raise ValueError("unknown source")
|
|
dist = {src: 0}
|
|
pq = [(0, src)]
|
|
while pq:
|
|
d, u = heapq.heappop(pq)
|
|
if d != dist.get(u, None):
|
|
continue
|
|
for v, w in adj.get(u, []):
|
|
if w < 0:
|
|
raise ValueError("negative weight")
|
|
nd = d + w
|
|
if nd < dist.get(v, float("inf")):
|
|
dist[v] = nd
|
|
heapq.heappush(pq, (nd, v))
|
|
# ensure all nodes present
|
|
for n in adj:
|
|
dist.setdefault(n, float("inf"))
|
|
return dist
|
|
|
|
|
|
if __name__ == "__main__":
|
|
g = {"a": [("b", 1), ("c", 4)], "b": [("c", 2)], "c": []}
|
|
print(dijkstra(g, "a"))
|