beam: BFS shortest path module with tests

wizard-work/beam/bfs_shortest_path.py

@@ -0,0 +1,66 @@
+"""BFS shortest path (unweighted) with deterministic path reconstruction.
+
+Public API:
+    shortest_path(adj, src, dst) -> list[int] | None
+        Returns the vertex sequence from src to dst along a shortest path,
+        or None if dst is unreachable. src==dst returns [src].
+    shortest_distances(adj, src) -> dict[int, int]
+        Distance in edges from src to every reachable vertex.
+"""
+
+from collections import deque
+
+
+def shortest_distances(adj, src):
+    if src not in adj:
+        # isolated source still reachable to itself
+        pass
+    dist = {src: 0}
+    q = deque([src])
+    while q:
+        u = q.popleft()
+        for v in adj.get(u, ()):
+            if v not in dist:
+                dist[v] = dist[u] + 1
+                q.append(v)
+    return dist
+
+
+def shortest_path(adj, src, dst):
+    if src == dst:
+        if src in adj or src in _all_vertices(adj):
+            return [src]
+        # allow unknown isolated vertex only if explicitly present as key? We require vertex existence.
+        return None
+    if src not in _all_vertices(adj) or dst not in _all_vertices(adj):
+        return None
+    parent = {src: None}
+    q = deque([src])
+    found = False
+    while q:
+        u = q.popleft()
+        if u == dst:
+            found = True
+            break
+        for v in adj.get(u, ()):
+            if v not in parent:
+                parent[v] = u
+                q.append(v)
+    if dst not in parent:
+        return None
+    # reconstruct
+    path = []
+    cur = dst
+    while cur is not None:
+        path.append(cur)
+        cur = parent[cur]
+    path.reverse()
+    return path
+
+
+def _all_vertices(adj):
+    s = set(adj.keys())
+    for u, nbrs in adj.items():
+        for v in nbrs:
+            s.add(v)
+    return s

wizard-work/beam/test_bfs.py

@@ -0,0 +1,39 @@
+import sys
+sys.path.insert(0, "/root/.hermes/wizard-outputs/beam")
+from bfs_shortest_path import shortest_path, shortest_distances
+
+def test_basic():
+    adj = {0:[1,2],1:[3],2:[3],3:[]}
+    assert shortest_path(adj,0,3) in ([0,1,3],[0,2,3])
+    assert shortest_distances(adj,0) == {0:0,1:1,2:1,3:2}
+
+def test_same():
+    adj = {5:[]}
+    assert shortest_path(adj,5,5) == [5]
+
+def test_unreachable():
+    adj = {0:[1],2:[]}
+    assert shortest_path(adj,0,2) is None
+
+def test_unknown_vertex():
+    adj = {0:[1]}
+    assert shortest_path(adj,0,99) is None
+    assert shortest_path(adj,99,0) is None
+
+def test_empty_graph():
+    adj = {}
+    assert shortest_path(adj,1,1) is None  # vertex not present
+    assert shortest_distances(adj,1) == {1:0}  # isolated allowed for distances? spec: yes if src given
+
+def test_directed():
+    adj = {0:[1],1:[]}
+    assert shortest_path(adj,1,0) is None
+
+def test_path_reconstruction_long():
+    adj = {i:[i+1] for i in range(5)}
+    assert shortest_path(adj,0,5) == [0,1,2,3,4,5]
+
+if __name__ == "__main__":
+    for name in [n for n in globals() if n.startswith("test_")]:
+        globals()[name]()
+    print("ALL TESTS PASSED")