Topological sort in Python
Produces a topological ordering by reversing DFS finish order.
- File
topological_sort.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/topological-sort/. The authors are listed in prov/KDXVF9.jsonl.
WHITE, GRAY, BLACK = 0, 1, 2
def topological_sort(adj: list[list[int]]) -> list[int] | None:
"""Return a topological ordering of the directed graph, or None if it has a cycle.
The graph has the vertices 0 .. n-1 and the adjacency lists adj. The order is the
reverse postorder of a DFS that tries the roots in the order 0 .. n-1 and each list
in its given order, so it is the vertices by decreasing DFS finish time. The
recursion is simulated with an explicit stack of (vertex, next index) pairs.
Time O(n + m), extra space O(n).
"""
n = len(adj)
color = [WHITE] * n # @line L1
post: list[int] = [] # @line L2: the list L, built from the back and reversed at the end
for root in range(n): # @line L3
if color[root] != WHITE: # @line L4
continue
color[root] = GRAY # @line L6
stack = [(root, 0)]
while stack:
u, i = stack[-1]
if i < len(adj[u]): # @line L7
stack[-1] = (u, i + 1)
v = adj[u][i]
if color[v] == GRAY: # @line L8
return None
if color[v] == WHITE: # @line L9
color[v] = GRAY # the call VISIT(v) starts with line L6
stack.append((v, 0))
else:
stack.pop()
color[u] = BLACK # @line L10
post.append(u) # @line L11
post.reverse()
return post # @line L5
test_topological_sort.py 105 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/topological-sort/. The authors are listed in prov/KDXVF9.jsonl.
import random
from topological_sort import topological_sort
def is_topological(adj: list[list[int]], order: list[int]) -> bool:
n = len(adj)
if sorted(order) != list(range(n)):
return False
pos = [0] * n
for i, v in enumerate(order):
pos[v] = i
return all(pos[u] < pos[v] for u in range(n) for v in adj[u])
def has_cycle(adj: list[list[int]]) -> bool:
"""Independent oracle: some vertex reaches itself by a nonempty path."""
n = len(adj)
for s in range(n):
seen = [False] * n
todo = list(adj[s])
while todo:
u = todo.pop()
if u == s:
return True
if not seen[u]:
seen[u] = True
todo.extend(adj[u])
return False
def reverse_postorder(adj: list[list[int]]) -> list[int]:
"""The textbook recursive DFS, used as the oracle for the exact order."""
n = len(adj)
seen = [False] * n
out: list[int] = []
def visit(u: int) -> None:
seen[u] = True
for v in adj[u]:
if not seen[v]:
visit(v)
out.append(u)
for u in range(n):
if not seen[u]:
visit(u)
return out[::-1]
def test_example() -> None:
adj = [[1, 3, 2], [2], [], [2, 4], [], [4, 6], []]
assert topological_sort(adj) == [5, 6, 0, 3, 4, 1, 2]
def test_example_with_cycle() -> None:
adj = [[1, 3, 2], [2], [0], [2, 4], [], [4, 6], []]
assert topological_sort(adj) is None
def test_small_cases() -> None:
assert topological_sort([]) == []
assert topological_sort([[]]) == [0]
assert topological_sort([[0]]) is None
assert topological_sort([[], [], []]) == [2, 1, 0]
assert topological_sort([[1], [0]]) is None
diamond = [[1, 2], [3], [3], []]
assert topological_sort(diamond) == [0, 2, 1, 3]
assert topological_sort([[1], [], [1]]) == [2, 0, 1]
def test_long_path() -> None:
n = 100000
adj = [[i + 1] for i in range(n - 1)] + [[]]
assert topological_sort(adj) == list(range(n))
adj[-1] = [0]
assert topological_sort(adj) is None
def test_random_against_oracles() -> None:
rng = random.Random(5150)
for trial in range(600):
n = rng.randint(0, 9)
p = rng.random()
if trial % 2 == 0:
perm = list(range(n))
rng.shuffle(perm)
adj = [[] for _ in range(n)]
for i in range(n):
for j in range(i + 1, n):
if rng.random() < p:
adj[perm[i]].append(perm[j])
else:
adj = [[v for v in range(n) if rng.random() < p / 2] for _ in range(n)]
for lst in adj:
rng.shuffle(lst)
got = topological_sort(adj)
if has_cycle(adj):
assert got is None
else:
assert got is not None and is_topological(adj, got)
assert got == reverse_postorder(adj)
What the next level needs
- the tests did not run in this build