beam: add bellman_ford with negative-cycle detection + tests

wizard-work/beam/bellman_ford.py

@@ -0,0 +1,84 @@
+"""Bellman-Ford shortest paths with negative-cycle detection.
+
+Given a directed weighted graph (possibly with negative edge weights) and a
+source vertex, compute shortest-path distances from the source. Detects
+negative-weight cycles reachable from the source.
+
+Returns (dist, parent, has_negative_cycle). dist maps vertex -> distance (inf
+if unreachable). parent maps vertex -> predecessor on a shortest path (or
+None for source/unreachable). has_negative_cycle is True iff a negative cycle
+reachable from the source exists.
+"""
+
+from __future__ import annotations
+from typing import Dict, Hashable, List, Optional, Tuple
+import math
+
+Vertex = Hashable
+Edge = Tuple[Vertex, Vertex, float]  # (u, v, weight)
+
+
+def bellman_ford(edges: List[Edge], n: int, source: Vertex,
+                 vertices: Optional[List[Vertex]] = None) -> Tuple[Dict, Dict, bool]:
+    if vertices is None:
+        # infer vertices 0..n-1 unless source not in range
+        vertices = list(range(n))
+    INF = float("inf")
+    dist: Dict[Vertex, float] = {v: INF for v in vertices}
+    parent: Dict[Vertex, Optional[Vertex]] = {v: None for v in vertices}
+    if source not in dist:
+        dist[source] = 0.0
+        parent[source] = None
+        if source not in vertices:
+            vertices = vertices + [source]
+            # ensure dist/parent cover source
+            if source not in dist:
+                dist[source] = 0.0
+                parent[source] = None
+    dist[source] = 0.0
+
+    # Relax all edges V-1 times.
+    V = len(vertices)
+    for _ in range(V - 1):
+        updated = False
+        for u, v, w in edges:
+            if dist.get(u, INF) != INF and dist[u] + w < dist.get(v, INF):
+                dist[v] = dist[u] + w
+                parent[v] = u
+                updated = True
+        if not updated:
+            break
+
+    # Detect negative cycle reachable from source.
+    has_neg = False
+    for u, v, w in edges:
+        if dist.get(u, INF) != INF and dist[u] + w < dist.get(v, INF):
+            has_neg = True
+            break
+    return dist, parent, has_neg
+
+
+def reconstruct_path(parent: Dict, source: Vertex, target: Vertex) -> Optional[List[Vertex]]:
+    if target not in parent:
+        return None
+    # walk
+    path = []
+    cur: Optional[Vertex] = target
+    seen = set()
+    while cur is not None:
+        if cur in seen:
+            return None  # cycle
+        seen.add(cur)
+        path.append(cur)
+        if cur == source:
+            break
+        cur = parent.get(cur)
+        if cur is None and path[-1] != source:
+            # unreachable
+            if target != source and parent.get(target) is None and target != source:
+                return None
+            break
+    if path and path[-1] != source:
+        return None
+    path.reverse()
+    return path

wizard-work/beam/test_bellman_ford.py

@@ -0,0 +1,57 @@
+import sys, os, math, unittest
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from bellman_ford import bellman_ford, reconstruct_path
+
+
+class TestBellmanFord(unittest.TestCase):
+    def test_simple(self):
+        edges = [(0,1,4),(0,2,1),(2,1,2),(1,3,1),(2,3,5)]
+        dist, parent, neg = bellman_ford(edges, 4, 0)
+        self.assertFalse(neg)
+        self.assertEqual(dist[0], 0)
+        self.assertEqual(dist[2], 1)
+        self.assertEqual(dist[1], 3)
+        self.assertEqual(dist[3], 4)
+
+    def test_negative_edges(self):
+        edges = [(0,1,4),(0,2,5),(1,2,-3),(2,3,4),(1,3,6)]
+        dist, parent, neg = bellman_ford(edges, 4, 0)
+        self.assertFalse(neg)
+        self.assertEqual(dist[2], 1)  # 0->1->2 = 4-3=1
+        self.assertEqual(dist[3], 5)  # 0->1->2->3 = 1+4=5
+
+    def test_negative_cycle(self):
+        edges = [(0,1,1),(1,2,-1),(2,0,-1)]  # cycle sum -1
+        dist, parent, neg = bellman_ford(edges, 3, 0)
+        self.assertTrue(neg)
+
+    def test_no_negative_cycle_unreachable(self):
+        edges = [(0,1,2),(2,3,-5),(3,2,1)]  # cycle in comp 2 not reachable
+        dist, parent, neg = bellman_ford(edges, 4, 0)
+        self.assertFalse(neg)
+        self.assertEqual(dist[2], float("inf"))
+
+    def test_path_reconstruct(self):
+        edges = [(0,1,4),(0,2,1),(2,1,2)]
+        dist, parent, neg = bellman_ford(edges, 3, 0)
+        self.assertEqual(reconstruct_path(parent, 0, 1), [0,2,1])
+
+    def test_single_vertex(self):
+        dist, parent, neg = bellman_ford([], 1, 0)
+        self.assertFalse(neg)
+        self.assertEqual(dist[0], 0)
+
+    def test_unreachable(self):
+        edges = [(0,1,5)]
+        dist, parent, neg = bellman_ford(edges, 3, 0)
+        self.assertEqual(dist[2], float("inf"))
+
+    def test_zero_weight(self):
+        edges = [(0,1,0),(1,2,0)]
+        dist, parent, neg = bellman_ford(edges, 3, 0)
+        self.assertEqual(dist[2], 0)
+        self.assertFalse(neg)
+
+
+if __name__ == "__main__":
+    unittest.main()