diff --git a/wizard-work/beam/dijkstra.py b/wizard-work/beam/dijkstra.py new file mode 100644 index 0000000..5d6e3e2 --- /dev/null +++ b/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")) diff --git a/wizard-work/beam/test_dijkstra.py b/wizard-work/beam/test_dijkstra.py new file mode 100644 index 0000000..ebc9809 --- /dev/null +++ b/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")