diff --git a/algorithms/graph/bellman_ford.py b/algorithms/graph/bellman_ford.py index 554074960..80772a994 100644 --- a/algorithms/graph/bellman_ford.py +++ b/algorithms/graph/bellman_ford.py @@ -37,7 +37,7 @@ def bellman_ford(graph: dict[str, dict[str, float]], source: str) -> bool: _initialize_single_source(graph, source, distance, predecessor) - num_vertices = len(graph) + num_vertices = len(distance) for _ in range(1, num_vertices): for current_node in graph: for neighbor in graph[current_node]: @@ -88,7 +88,7 @@ def bellman_ford_paths( _initialize_single_source(graph, source, distance, predecessor) - num_vertices = len(graph) + num_vertices = len(distance) for _ in range(1, num_vertices): for current_node in graph: if distance[current_node] == float("inf"): diff --git a/tests/test_graph.py b/tests/test_graph.py index 057544254..c4d23ae56 100644 --- a/tests/test_graph.py +++ b/tests/test_graph.py @@ -32,6 +32,7 @@ topological_sort, walls_and_gates, ) +from algorithms.graph.bellman_ford import bellman_ford_paths from algorithms.graph.dijkstra import Dijkstra @@ -233,6 +234,47 @@ def test_bellman_ford(self): } self.assertEqual(True, bellman_ford(graph2, "a")) + def test_single_edge_to_neighbor_only_vertex(self): + graph = {"a": {"b": 1}} + + self.assertTrue(bellman_ford(graph, "a")) + self.assertEqual( + bellman_ford_paths(graph, "a"), + ({"a": 0, "b": 1}, {"a": None, "b": "a"}, False), + ) + + def test_neighbor_only_sink_independent_of_vertex_order(self): + graphs = ( + {"a": {"b": 1}, "b": {"c": 1}}, + {"b": {"c": 1}, "a": {"b": 1}}, + ) + for graph in graphs: + with self.subTest(graph=graph): + self.assertTrue(bellman_ford(graph, "a")) + self.assertEqual( + bellman_ford_paths(graph, "a"), + ( + {"a": 0, "b": 1, "c": 2}, + {"a": None, "b": "a", "c": "b"}, + False, + ), + ) + + def test_negative_edge_to_neighbor_only_vertex(self): + graph = {"a": {"b": -1}} + + self.assertTrue(bellman_ford(graph, "a")) + self.assertEqual( + bellman_ford_paths(graph, "a"), + ({"a": 0, "b": -1}, {"a": None, "b": "a"}, False), + ) + + def test_negative_cycle_with_neighbor_only_sink(self): + graph = {"a": {"b": 1}, "b": {"a": -2, "c": 1}} + + self.assertFalse(bellman_ford(graph, "a")) + self.assertEqual(bellman_ford_paths(graph, "a"), (None, None, True)) + class TestConnectedComponentInGraph(unittest.TestCase): """