Yen's algorithm in Python
Finds the k shortest loopless paths by repeated spur-path computations with Dijkstra.
- File
yen_k_shortest_paths.py, 79 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/yen-k-shortest-paths/. The authors are listed in prov/0RC5G2.jsonl.
import heapq
def _dijkstra_path(
best: list[dict[int, float]], s: int, t: int, banned_vertices: set[int], banned_arcs: set[tuple[int, int]]
) -> tuple[float, list[int]] | None:
"""Shortest s-t path avoiding the banned vertices and arcs; weights must be nonnegative."""
dist: dict[int, float] = {s: 0}
parent: dict[int, int] = {}
heap: list[tuple[float, int]] = [(0, s)]
done: set[int] = set()
while heap:
d, u = heapq.heappop(heap)
if u in done:
continue
done.add(u)
if u == t:
path = [t]
while path[-1] != s:
path.append(parent[path[-1]])
path.reverse()
return d, path
for v, w in best[u].items():
if v in banned_vertices or (u, v) in banned_arcs or v in done:
continue
nd = d + w
if v not in dist or nd < dist[v]:
dist[v] = nd
parent[v] = u
heapq.heappush(heap, (nd, v))
return None
def yen_k_shortest_paths(
n: int, edges: list[tuple[int, int, float]], s: int, t: int, k: int
) -> list[tuple[float, list[int]]]:
"""The k shortest simple s-t paths, as (weight, vertex list), in nondecreasing order of weight.
Vertices are 0 .. n-1 and edges are (u, v, weight) with nonnegative weights. Parallel arcs are
merged to the lightest one, since a path is given by its vertices; self-loops are ignored,
because no simple path uses one. Fewer than k paths are returned when fewer exist.
Time O(k n (m + n log n)) with a binary heap.
"""
best: list[dict[int, float]] = [{} for _ in range(n)]
for u, v, w in edges:
if w < 0:
raise ValueError("Yen's algorithm with Dijkstra needs nonnegative weights")
if u != v and (v not in best[u] or w < best[u][v]):
best[u][v] = w
if k <= 0:
return []
first = _dijkstra_path(best, s, t, set(), set()) # @line L1
if first is None:
return []
a = [first]
candidates: list[tuple[float, list[int]]] = [] # @line L2: a min-heap of (weight, path)
seen = {tuple(first[1])}
while len(a) < k: # @line L3
_, prev = a[-1]
for i in range(len(prev) - 1): # @line L4
spur = prev[i] # @line L5
root = prev[: i + 1]
root_weight = sum(best[root[j]][root[j + 1]] for j in range(i))
banned_arcs = {(spur, p[i + 1]) for _, p in a if len(p) > i + 1 and p[: i + 1] == root} # @line L6
banned_vertices = set(root[:-1]) # @line L7
found = _dijkstra_path(best, spur, t, banned_vertices, banned_arcs) # @line L8
if found is not None: # @line L9
path = root[:-1] + found[1]
if tuple(path) not in seen:
seen.add(tuple(path))
heapq.heappush(candidates, (root_weight + found[0], path))
# @line L10: the banned sets are local, so nothing has to be restored
if not candidates: # @line L11
break
a.append(heapq.heappop(candidates)) # @line L12
return a # @line L13
test_yen_k_shortest_paths.py 89 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/yen-k-shortest-paths/. The authors are listed in prov/0RC5G2.jsonl.
import random
import unittest
from yen_k_shortest_paths import yen_k_shortest_paths
NAMES = "CDEFGH"
EXAMPLE = [
("C", "D", 3), ("C", "E", 2), ("D", "F", 4), ("E", "D", 1), ("E", "F", 2),
("E", "G", 3), ("F", "G", 2), ("F", "H", 1), ("G", "H", 2),
]
def named(edges):
return [(NAMES.index(a), NAMES.index(b), w) for a, b, w in edges]
def all_simple_paths(n, edges, s, t):
"""Independent oracle: every simple s-t path by depth-first search, with its least weight."""
best = {}
for u, v, w in edges:
if u != v and ((u, v) not in best or w < best[(u, v)]):
best[(u, v)] = w
out = [[] for _ in range(n)]
for (u, v), w in best.items():
out[u].append((v, w))
found = []
def go(u, path, weight):
if u == t:
found.append((weight, tuple(path)))
return
for v, w in out[u]:
if v not in path:
path.append(v)
go(v, path, weight + w)
path.pop()
go(s, [s], 0)
return sorted(found), best
class TestYen(unittest.TestCase):
def test_example(self):
got = yen_k_shortest_paths(6, named(EXAMPLE), 0, 5, 3)
self.assertEqual([(w, "".join(NAMES[x] for x in p)) for w, p in got], [(5, "CEFH"), (7, "CEGH"), (8, "CDFH")])
def test_all_paths_of_example(self):
got = yen_k_shortest_paths(6, named(EXAMPLE), 0, 5, 100)
self.assertEqual([w for w, _ in got], [5, 7, 8, 8, 8, 11, 11])
def test_unreachable_and_trivial(self):
self.assertEqual(yen_k_shortest_paths(3, [(0, 1, 1)], 0, 2, 4), [])
self.assertEqual(yen_k_shortest_paths(2, [(0, 1, 1)], 0, 0, 3), [(0, [0])])
self.assertEqual(yen_k_shortest_paths(2, [(0, 1, 1)], 0, 1, 0), [])
def test_negative_weights_and_cycles_rejected(self):
with self.assertRaises(ValueError):
yen_k_shortest_paths(2, [(0, 1, -1)], 0, 1, 2)
with self.assertRaises(ValueError):
yen_k_shortest_paths(3, [(0, 1, 1), (1, 2, 1), (2, 1, -5)], 0, 2, 2)
def test_cycles_and_parallel_arcs(self):
edges = [(0, 1, 1), (1, 0, 1), (1, 2, 1), (0, 2, 5), (0, 2, 3), (2, 2, 0)]
self.assertEqual(yen_k_shortest_paths(3, edges, 0, 2, 5), [(2, [0, 1, 2]), (3, [0, 2])])
def test_random_against_enumeration(self):
rng = random.Random(1971)
for _ in range(600):
n = rng.randint(1, 7)
m = rng.randint(0, 3 * n)
edges = [(rng.randrange(n), rng.randrange(n), rng.randint(0, 6)) for _ in range(m)]
s, t = rng.randrange(n), rng.randrange(n)
k = rng.randint(1, 12)
expected, best = all_simple_paths(n, edges, s, t)
got = yen_k_shortest_paths(n, edges, s, t, k)
self.assertEqual([w for w, _ in got], [w for w, _ in expected[:k]])
paths = [tuple(p) for _, p in got]
self.assertEqual(len(set(paths)), len(paths))
for w, p in got:
self.assertEqual((p[0], p[-1]), (s, t))
self.assertEqual(len(set(p)), len(p))
self.assertEqual(sum(best[(p[i], p[i + 1])] for i in range(len(p) - 1)), w)
if __name__ == "__main__":
unittest.main()
What the next level needs
- the tests did not run in this build