Warshall's algorithm in Python
A dynamic program over intermediate vertices computes the transitive closure in cubic time.
- File
warshall_algorithm.py, 27 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/warshall-algorithm/. The authors are listed in prov/036P3Q.jsonl.
def warshall_algorithm(adj: list[list[int]]) -> list[list[bool]]:
"""Return the transitive closure R of a directed graph as an n x n Boolean matrix.
The graph has vertices 0 .. n-1 and adjacency lists adj. Parallel arcs (repeated
entries) and self-loops (v in adj[v]) are allowed. R[u][v] is True exactly when
there is a walk of length at least 1 from u to v, so R[v][v] is True exactly
when v lies on a cycle (a loop counts). Each row is kept as a Python int used as
a bit vector, so line L5 is one OR of two rows.
Time Theta(n^3) bit operations (n^3 / w word operations), space Theta(n^2) bits.
"""
n = len(adj)
row = [0] * n # @line L1
for u in range(n):
for v in adj[u]:
row[u] |= 1 << v
for k in range(n): # @line L2
bit = 1 << k
rk = row[k]
for i in range(n): # @line L3
if row[i] & bit: # @line L4
row[i] |= rk # @line L5
# row[k] itself can change in pass k only by OR with itself, so rk stays valid.
return [[bool(row[u] >> v & 1) for v in range(n)] for u in range(n)] # @line L6
test_warshall_algorithm.py 62 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/warshall-algorithm/. The authors are listed in prov/036P3Q.jsonl.
import random
from warshall_algorithm import warshall_algorithm
def oracle(adj: list[list[int]]) -> list[list[bool]]:
"""R[u][v] by a search from the successors of u (walks of length at least 1)."""
n = len(adj)
out = []
for u in range(n):
seen = set(adj[u])
todo = list(seen)
while todo:
x = todo.pop()
for y in adj[x]:
if y not in seen:
seen.add(y)
todo.append(y)
out.append([v in seen for v in range(n)])
return out
def test_example() -> None:
adj = [[1, 3, 2], [2], [0], [2, 4], [], [4, 6], [5]]
r = warshall_algorithm(adj)
a = [True] * 5 + [False] * 2
assert r[:4] == [a] * 4
assert r[4] == [False] * 7
assert r[5] == r[6] == [False] * 4 + [True] * 3
assert sum(map(sum, r)) == 26
def test_edges() -> None:
assert warshall_algorithm([]) == []
assert warshall_algorithm([[]]) == [[False]]
assert warshall_algorithm([[0]]) == [[True]]
assert warshall_algorithm([[1, 1], []]) == [[False, True], [False, False]]
# 3-cycle listed so that the wrong loop order (i outside k) misses R[0][0]
assert warshall_algorithm([[2], [0], [1]]) == [[True] * 3] * 3
# undirected path 0 - 1 - 2 given with both arcs: one component, full matrix
assert warshall_algorithm([[1], [0, 2], [1]]) == [[True] * 3] * 3
def test_path() -> None:
n = 40
adj = [[i + 1] for i in range(n - 1)] + [[]]
r = warshall_algorithm(adj)
assert all(r[u][v] == (u < v) for u in range(n) for v in range(n))
def test_random() -> None:
rng = random.Random(36)
for _ in range(700):
n = rng.randrange(1, 10)
m = rng.randrange(0, 3 * n)
adj: list[list[int]] = [[] for _ in range(n)]
for _ in range(m):
adj[rng.randrange(n)].append(rng.randrange(n))
assert warshall_algorithm(adj) == oracle(adj)
What the next level needs
- the tests did not run in this build