Skip to content

Commit a9f2e72

Browse files
sangampaudel530poyeacclaussCopilot
authored
Added Johnson's algorithm for all-pairs shortest paths (#13340)
* Fix typos in Johnson's algorithm (nd -> and) to pass codespell * Rename type aliases and h parameter to follow snake_case and descriptive naming * Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: John Law <johnlaw.po@gmail.com> Co-authored-by: Christian Clauss <cclauss@me.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
1 parent 33a8e0f commit a9f2e72

2 files changed

Lines changed: 142 additions & 0 deletions

File tree

graphs/johnson.py

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
import heapq
2+
from collections.abc import Hashable
3+
4+
Node = Hashable
5+
edge = tuple[Node, Node, float]
6+
adjacency = dict[Node, list[tuple[Node, float]]]
7+
8+
9+
def _collect_nodes_and_edges(graph: adjacency) -> tuple[list[Node], list[edge]]:
10+
nodes = set()
11+
edges: list[edge] = []
12+
for u, neighbors in graph.items():
13+
nodes.add(u)
14+
for v, w in neighbors:
15+
nodes.add(v)
16+
edges.append((u, v, w))
17+
return list(nodes), edges
18+
19+
20+
def _bellman_ford(nodes: list[Node], edges: list[edge]) -> dict[Node, float]:
21+
"""
22+
Bellman-Ford relaxation to compute potentials h[v] for all vertices.
23+
Raises ValueError if a negative weight cycle exists.
24+
"""
25+
dist: dict[Node, float] = dict.fromkeys(nodes, 0.0)
26+
n = len(nodes)
27+
28+
for _ in range(n - 1):
29+
updated = False
30+
for u, v, w in edges:
31+
if dist[u] + w < dist[v]:
32+
dist[v] = dist[u] + w
33+
updated = True
34+
if not updated:
35+
break
36+
else:
37+
for u, v, w in edges:
38+
if dist[u] + w < dist[v]:
39+
raise ValueError("Negative weight cycle detected")
40+
return dist
41+
42+
43+
def _dijkstra(
44+
start: Node,
45+
nodes: list[Node],
46+
graph: adjacency,
47+
potentials: dict[Node, float],
48+
) -> dict[Node, float]:
49+
"""
50+
Dijkstra over reweighted graph, using potentials h to make weights non-negative.
51+
Returns distances from start in the reweighted space.
52+
"""
53+
inf = float("inf")
54+
dist: dict[Node, float] = dict.fromkeys(nodes, inf)
55+
dist[start] = 0.0
56+
heap: list[tuple[float, Node]] = [(0.0, start)]
57+
58+
while heap:
59+
d_u, u = heapq.heappop(heap)
60+
if d_u > dist[u]:
61+
continue
62+
for v, w in graph.get(u, []):
63+
w_prime = w + potentials[u] - potentials[v]
64+
if w_prime < 0:
65+
raise ValueError(
66+
"Negative edge weight after reweighting: numeric error"
67+
)
68+
new_dist = d_u + w_prime
69+
if new_dist < dist[v]:
70+
dist[v] = new_dist
71+
heapq.heappush(heap, (new_dist, v))
72+
return dist
73+
74+
75+
def johnson(graph: adjacency) -> dict[Node, dict[Node, float]]:
76+
"""
77+
Compute all-pairs shortest paths using Johnson's algorithm.
78+
79+
Reference:
80+
https://en.wikipedia.org/wiki/Johnson%27s_algorithm
81+
82+
Args:
83+
graph: adjacency list {u: [(v, weight), ...], ...}
84+
85+
Returns:
86+
dict of dicts: dist[u][v] = shortest distance from u to v
87+
88+
Raises:
89+
ValueError: if a negative weight cycle is detected
90+
91+
Example:
92+
>>> g = {
93+
... 0: [(1, 3), (2, 8), (4, -4)],
94+
... 1: [(3, 1), (4, 7)],
95+
... 2: [(1, 4)],
96+
... 3: [(0, 2), (2, -5)],
97+
... 4: [(3, 6)],
98+
... }
99+
>>> round(johnson(g)[0][3], 2)
100+
2.0
101+
"""
102+
nodes, edges = _collect_nodes_and_edges(graph)
103+
potentials = _bellman_ford(nodes, edges)
104+
105+
all_pairs: dict[Node, dict[Node, float]] = {}
106+
inf = float("inf")
107+
for s in nodes:
108+
dist_reweighted = _dijkstra(s, nodes, graph, potentials)
109+
dists_orig: dict[Node, float] = {}
110+
for v in nodes:
111+
d_prime = dist_reweighted[v]
112+
if d_prime < inf:
113+
dists_orig[v] = d_prime - potentials[s] + potentials[v]
114+
else:
115+
dists_orig[v] = inf
116+
all_pairs[s] = dists_orig
117+
118+
return all_pairs

graphs/tests/test_johnson.py

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
import math
2+
3+
import pytest
4+
5+
from graphs.johnson import johnson
6+
7+
8+
def test_johnson_basic():
9+
g = {
10+
0: [(1, 3), (2, 8), (4, -4)],
11+
1: [(3, 1), (4, 7)],
12+
2: [(1, 4)],
13+
3: [(0, 2), (2, -5)],
14+
4: [(3, 6)],
15+
}
16+
dist = johnson(g)
17+
assert math.isclose(dist[0][3], 2.0, abs_tol=1e-9)
18+
assert math.isclose(dist[3][2], -5.0, abs_tol=1e-9)
19+
20+
21+
def test_johnson_negative_cycle():
22+
g2 = {0: [(1, 1)], 1: [(0, -3)]}
23+
with pytest.raises(ValueError):
24+
johnson(g2)

0 commit comments

Comments
 (0)