Skip to content

Commit 93d6514

Browse files
Fix typos in Johnson's algorithm (nd -> and) to pass codespell
1 parent 788d95b commit 93d6514

2 files changed

Lines changed: 140 additions & 0 deletions

File tree

graphs/johnson.py

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
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+
# One more iteration to check for negative cycles
38+
for u, v, w in edges:
39+
if dist[u] + w < dist[v]:
40+
raise ValueError("Negative weight cycle detected")
41+
return dist
42+
43+
44+
def _dijkstra(
45+
start: Node,
46+
nodes: list[Node],
47+
graph: Adjacency,
48+
h: dict[Node, float],
49+
) -> dict[Node, float]:
50+
"""
51+
Dijkstra over reweighted graph, using potentials h to make weights non-negative.
52+
Returns distances from start in the reweighted space.
53+
"""
54+
inf = float("inf")
55+
dist: dict[Node, float] = dict.fromkeys(nodes, inf)
56+
dist[start] = 0.0
57+
heap: list[tuple[float, Node]] = [(0.0, start)]
58+
59+
while heap:
60+
d_u, u = heapq.heappop(heap)
61+
if d_u > dist[u]:
62+
continue
63+
for v, w in graph.get(u, []):
64+
w_prime = w + h[u] - h[v]
65+
if w_prime < 0:
66+
raise ValueError(
67+
"Negative edge weight after reweighting: numeric error"
68+
)
69+
new_dist = d_u + w_prime
70+
if new_dist < dist[v]:
71+
dist[v] = new_dist
72+
heapq.heappush(heap, (new_dist, v))
73+
return dist
74+
75+
76+
def johnson(graph: Adjacency) -> dict[Node, dict[Node, float]]:
77+
"""
78+
Compute all-pairs shortest paths using Johnson's algorithm.
79+
80+
Args:
81+
graph: adjacency list {u: [(v, weight), ...], ...}
82+
83+
Returns:
84+
dict of dicts: dist[u][v] = shortest distance from u to v
85+
86+
Raises:
87+
ValueError: if a negative weight cycle is detected
88+
89+
Example:
90+
>>> g = {
91+
... 0: [(1, 3), (2, 8), (4, -4)],
92+
... 1: [(3, 1), (4, 7)],
93+
... 2: [(1, 4)],
94+
... 3: [(0, 2), (2, -5)],
95+
... 4: [(3, 6)],
96+
... }
97+
>>> round(johnson(g)[0][3], 2)
98+
2.0
99+
"""
100+
nodes, edges = _collect_nodes_and_edges(graph)
101+
h = _bellman_ford(nodes, edges)
102+
103+
all_pairs: dict[Node, dict[Node, float]] = {}
104+
inf = float("inf")
105+
for s in nodes:
106+
dist_reweighted = _dijkstra(s, nodes, graph, h)
107+
dists_orig: dict[Node, float] = {}
108+
for v in nodes:
109+
d_prime = dist_reweighted[v]
110+
if d_prime < inf:
111+
dists_orig[v] = d_prime - h[s] + h[v]
112+
else:
113+
dists_orig[v] = inf
114+
all_pairs[s] = dists_orig
115+
116+
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)