Dijkstra's algorithm in Python
Greedily finalizes the closest unfinalized vertex using a priority queue; requires nonnegative weights.
- File
dijkstra.py, 39 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/dijkstra/. The authors are listed in prov/27XR9T.jsonl.
import heapq
import math
def dijkstra(
adj: list[list[tuple[int, float]]], s: int
) -> tuple[list[float], list[int | None]]:
"""Return (d, parent): shortest-path distances from s and a shortest-path tree.
The graph has the vertices 0 .. n-1; adj[u] lists the edges (v, w) out of u.
Precondition: every weight w is nonnegative (a ValueError is raised otherwise).
d[v] is math.inf when v is not reachable from s, and parent[v] is then None.
The queue is a binary heap with lazy deletion: an entry whose vertex is already
done is skipped. Time O(n + m log m), which is O((n + m) log n) without parallel
edges; extra space O(n + m).
"""
n = len(adj)
for u in range(n):
for _, w in adj[u]:
if w < 0:
raise ValueError("Dijkstra's algorithm needs nonnegative weights")
d = [math.inf] * n # @line L1
parent: list[int | None] = [None] * n
done = [False] * n
d[s] = 0 # @line L2
q: list[tuple[float, int]] = [(0, s)] # @line L3
while q: # @line L4
_, u = heapq.heappop(q) # @line L5
if not done[u]: # @line L6
done[u] = True # @line L7
for v, w in adj[u]: # @line L8
if d[u] + w < d[v]: # @line L9
d[v] = d[u] + w # @line L10
parent[v] = u
heapq.heappush(q, (d[v], v)) # @line L11
return d, parent # @line L12
test_dijkstra.py 102 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/dijkstra/. The authors are listed in prov/27XR9T.jsonl.
import math
import random
import pytest
from dijkstra import dijkstra
NAMES = "sabcdt"
EDGES = [("s", "a", 7), ("s", "b", 2), ("b", "a", 3), ("b", "c", 8), ("a", "c", 1),
("a", "d", 6), ("c", "d", 2), ("c", "t", 9), ("d", "t", 3)]
def build(n: int, edges: list[tuple[int, int, float]]) -> list[list[tuple[int, float]]]:
adj: list[list[tuple[int, float]]] = [[] for _ in range(n)]
for u, v, w in edges:
adj[u].append((v, w))
return adj
def oracle(n: int, edges: list[tuple[int, int, float]], s: int) -> list[float]:
"""Independent check: n - 1 rounds of relaxing every edge (Bellman-Ford)."""
d = [math.inf] * n
d[s] = 0
for _ in range(n - 1):
for u, v, w in edges:
if d[u] + w < d[v]:
d[v] = d[u] + w
return d
def check_tree(adj, s, d, parent) -> None:
for v in range(len(adj)):
if v == s:
assert parent[v] is None and d[v] == 0
elif d[v] == math.inf:
assert parent[v] is None
else:
u = parent[v]
assert u is not None
assert any(x == v and d[u] + w == d[v] for x, w in adj[u])
def test_example() -> None:
ix = {c: i for i, c in enumerate(NAMES)}
adj = build(6, [(ix[u], ix[v], w) for u, v, w in EDGES])
d, parent = dijkstra(adj, 0)
assert d == [0, 5, 2, 6, 8, 11]
assert parent == [None, 2, 0, 1, 3, 4]
def test_single_vertex_and_self_loop() -> None:
assert dijkstra([[]], 0) == ([0], [None])
assert dijkstra([[(0, 5)]], 0) == ([0], [None])
def test_unreachable() -> None:
adj = build(4, [(0, 1, 3), (2, 3, 1), (3, 1, 0)])
d, parent = dijkstra(adj, 0)
assert d == [0, 3, math.inf, math.inf]
assert parent == [None, 0, None, None]
def test_zero_weights_and_parallel_edges() -> None:
adj = build(3, [(0, 1, 0), (0, 1, 4), (1, 2, 0), (0, 2, 1)])
d, _ = dijkstra(adj, 0)
assert d == [0, 0, 0]
def test_negative_edge_rejected() -> None:
# The page's counterexample: without the check the answer for c would be 3, not 2.
adj = build(4, [(0, 1, 2), (1, 3, 1), (0, 2, 4), (2, 1, -3)])
with pytest.raises(ValueError):
dijkstra(adj, 0)
def test_negative_cycle_rejected() -> None:
adj = build(3, [(0, 1, 1), (1, 2, 2), (2, 1, -3)])
with pytest.raises(ValueError):
dijkstra(adj, 0)
def test_large_weights_exact() -> None:
big = 2**62
adj = build(3, [(0, 1, big), (1, 2, big), (0, 2, 2**63 + 1)])
d, _ = dijkstra(adj, 0)
assert d[2] == 2**63
def test_random_against_bellman_ford() -> None:
rng = random.Random(1959)
for _ in range(600):
n = rng.randint(1, 8)
edges = [(rng.randrange(n), rng.randrange(n), rng.randint(0, 9))
for _ in range(rng.randint(0, 3 * n))]
adj = build(n, edges)
s = rng.randrange(n)
d, parent = dijkstra(adj, s)
assert d == oracle(n, edges, s)
check_tree(adj, s, d, parent)
What the next level needs
- the tests did not run in this build