beam: dijkstra shortest path with tests

wizard-work/beam/dijkstra.py

@@ -0,0 +1,34 @@
+"""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"))

wizard-work/beam/test_dijkstra.py

@@ -0,0 +1,31 @@
+from dijkstra import dijkstra
+
+def test_basic():
+    g = {"a": [("b", 1), ("c", 4)], "b": [("c", 2)], "c": []}
+    d = dijkstra(g, "a")
+    assert d["a"] == 0 and d["b"] == 1 and d["c"] == 3
+
+def test_unreachable():
+    g = {"a": [], "b": []}
+    d = dijkstra(g, "a")
+    assert d["b"] == float("inf")
+
+def test_unknown_src():
+    try:
+        dijkstra({"a": []}, "z")
+        raise AssertionError("expected")
+    except ValueError:
+        pass
+
+def test_negative():
+    try:
+        dijkstra({"a": [("b", -1)]}, "a")
+        raise AssertionError("expected")
+    except ValueError:
+        pass
+
+if __name__ == "__main__":
+    for n in [x for x in globals() if x.startswith("test_")]:
+        globals()[n]()
+        print("PASS", n)
+    print("ALL PASS")