diff --git a/wizard-work/beam/bellman_ford.py b/wizard-work/beam/bellman_ford.py new file mode 100644 index 0000000..29f47c8 --- /dev/null +++ b/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 diff --git a/wizard-work/beam/test_bellman_ford.py b/wizard-work/beam/test_bellman_ford.py new file mode 100644 index 0000000..f7e818a --- /dev/null +++ b/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()