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>
This commit is contained in:
co-authored by
Copilot Autofix powered by AI
John Law
Christian Clauss
parent
33a8e0f21a
commit
a9f2e72541
@@ -0,0 +1,118 @@
|
||||
import heapq
|
||||
from collections.abc import Hashable
|
||||
|
||||
Node = Hashable
|
||||
edge = tuple[Node, Node, float]
|
||||
adjacency = dict[Node, list[tuple[Node, float]]]
|
||||
|
||||
|
||||
def _collect_nodes_and_edges(graph: adjacency) -> tuple[list[Node], list[edge]]:
|
||||
nodes = set()
|
||||
edges: list[edge] = []
|
||||
for u, neighbors in graph.items():
|
||||
nodes.add(u)
|
||||
for v, w in neighbors:
|
||||
nodes.add(v)
|
||||
edges.append((u, v, w))
|
||||
return list(nodes), edges
|
||||
|
||||
|
||||
def _bellman_ford(nodes: list[Node], edges: list[edge]) -> dict[Node, float]:
|
||||
"""
|
||||
Bellman-Ford relaxation to compute potentials h[v] for all vertices.
|
||||
Raises ValueError if a negative weight cycle exists.
|
||||
"""
|
||||
dist: dict[Node, float] = dict.fromkeys(nodes, 0.0)
|
||||
n = len(nodes)
|
||||
|
||||
for _ in range(n - 1):
|
||||
updated = False
|
||||
for u, v, w in edges:
|
||||
if dist[u] + w < dist[v]:
|
||||
dist[v] = dist[u] + w
|
||||
updated = True
|
||||
if not updated:
|
||||
break
|
||||
else:
|
||||
for u, v, w in edges:
|
||||
if dist[u] + w < dist[v]:
|
||||
raise ValueError("Negative weight cycle detected")
|
||||
return dist
|
||||
|
||||
|
||||
def _dijkstra(
|
||||
start: Node,
|
||||
nodes: list[Node],
|
||||
graph: adjacency,
|
||||
potentials: dict[Node, float],
|
||||
) -> dict[Node, float]:
|
||||
"""
|
||||
Dijkstra over reweighted graph, using potentials h to make weights non-negative.
|
||||
Returns distances from start in the reweighted space.
|
||||
"""
|
||||
inf = float("inf")
|
||||
dist: dict[Node, float] = dict.fromkeys(nodes, inf)
|
||||
dist[start] = 0.0
|
||||
heap: list[tuple[float, Node]] = [(0.0, start)]
|
||||
|
||||
while heap:
|
||||
d_u, u = heapq.heappop(heap)
|
||||
if d_u > dist[u]:
|
||||
continue
|
||||
for v, w in graph.get(u, []):
|
||||
w_prime = w + potentials[u] - potentials[v]
|
||||
if w_prime < 0:
|
||||
raise ValueError(
|
||||
"Negative edge weight after reweighting: numeric error"
|
||||
)
|
||||
new_dist = d_u + w_prime
|
||||
if new_dist < dist[v]:
|
||||
dist[v] = new_dist
|
||||
heapq.heappush(heap, (new_dist, v))
|
||||
return dist
|
||||
|
||||
|
||||
def johnson(graph: adjacency) -> dict[Node, dict[Node, float]]:
|
||||
"""
|
||||
Compute all-pairs shortest paths using Johnson's algorithm.
|
||||
|
||||
Reference:
|
||||
https://en.wikipedia.org/wiki/Johnson%27s_algorithm
|
||||
|
||||
Args:
|
||||
graph: adjacency list {u: [(v, weight), ...], ...}
|
||||
|
||||
Returns:
|
||||
dict of dicts: dist[u][v] = shortest distance from u to v
|
||||
|
||||
Raises:
|
||||
ValueError: if a negative weight cycle is detected
|
||||
|
||||
Example:
|
||||
>>> g = {
|
||||
... 0: [(1, 3), (2, 8), (4, -4)],
|
||||
... 1: [(3, 1), (4, 7)],
|
||||
... 2: [(1, 4)],
|
||||
... 3: [(0, 2), (2, -5)],
|
||||
... 4: [(3, 6)],
|
||||
... }
|
||||
>>> round(johnson(g)[0][3], 2)
|
||||
2.0
|
||||
"""
|
||||
nodes, edges = _collect_nodes_and_edges(graph)
|
||||
potentials = _bellman_ford(nodes, edges)
|
||||
|
||||
all_pairs: dict[Node, dict[Node, float]] = {}
|
||||
inf = float("inf")
|
||||
for s in nodes:
|
||||
dist_reweighted = _dijkstra(s, nodes, graph, potentials)
|
||||
dists_orig: dict[Node, float] = {}
|
||||
for v in nodes:
|
||||
d_prime = dist_reweighted[v]
|
||||
if d_prime < inf:
|
||||
dists_orig[v] = d_prime - potentials[s] + potentials[v]
|
||||
else:
|
||||
dists_orig[v] = inf
|
||||
all_pairs[s] = dists_orig
|
||||
|
||||
return all_pairs
|
||||
@@ -0,0 +1,24 @@
|
||||
import math
|
||||
|
||||
import pytest
|
||||
|
||||
from graphs.johnson import johnson
|
||||
|
||||
|
||||
def test_johnson_basic():
|
||||
g = {
|
||||
0: [(1, 3), (2, 8), (4, -4)],
|
||||
1: [(3, 1), (4, 7)],
|
||||
2: [(1, 4)],
|
||||
3: [(0, 2), (2, -5)],
|
||||
4: [(3, 6)],
|
||||
}
|
||||
dist = johnson(g)
|
||||
assert math.isclose(dist[0][3], 2.0, abs_tol=1e-9)
|
||||
assert math.isclose(dist[3][2], -5.0, abs_tol=1e-9)
|
||||
|
||||
|
||||
def test_johnson_negative_cycle():
|
||||
g2 = {0: [(1, 1)], 1: [(0, -3)]}
|
||||
with pytest.raises(ValueError):
|
||||
johnson(g2)
|
||||
Reference in New Issue
Block a user