Depth-first search in Python
Explores as deep as possible along each branch before backtracking, using a stack or recursion.
- File
depth_first_search.py, 46 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/depth-first-search/. The authors are listed in prov/Q8B7W6.jsonl.
WHITE, GRAY, BLACK = 0, 1, 2
def depth_first_search(adj: list[list[int]]) -> tuple[list[int], list[int], list[int | None]]:
"""Run a full DFS on the graph with vertices 0 .. n-1 and adjacency lists adj.
Returns (d, f, parent): the discovery time, the finish time and the DFS parent of
each vertex. The times are 1 .. 2n. The roots are tried in the order 0 .. n-1 and
each list in its given order, so the result is the one of the recursive procedure.
The recursion is simulated with an explicit stack of (vertex, next index) pairs,
so a long path does not overflow the Python stack. Time O(n + m), extra space O(n).
"""
n = len(adj)
color = [WHITE] * n # @line L1
parent: list[int | None] = [None] * n
d = [0] * n
f = [0] * n
time = 0 # @line L2
for root in range(n): # @line L3
if color[root] != WHITE: # @line L4
continue
# DFS-VISIT(root), with the call stack made explicit.
time += 1 # @line L5
d[root] = time
color[root] = GRAY
stack = [(root, 0)]
while stack:
u, i = stack[-1]
if i < len(adj[u]): # @line L6
stack[-1] = (u, i + 1)
v = adj[u][i]
if color[v] == WHITE: # @line L7
parent[v] = u # @line L8
time += 1 # @line L9: the call DFS-VISIT(v) starts with line L5
d[v] = time
color[v] = GRAY
stack.append((v, 0))
else:
stack.pop()
time += 1 # @line L10
f[u] = time
color[u] = BLACK
return d, f, parent
test_depth_first_search.py 86 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/depth-first-search/. The authors are listed in prov/Q8B7W6.jsonl.
import random
import sys
from depth_first_search import depth_first_search
def recursive_dfs(adj: list[list[int]]) -> tuple[list[int], list[int], list[int | None]]:
"""The textbook recursive procedure, used as the oracle on small graphs."""
n = len(adj)
seen = [False] * n
d, f = [0] * n, [0] * n
parent: list[int | None] = [None] * n
clock = [0]
def visit(u: int) -> None:
clock[0] += 1
d[u] = clock[0]
seen[u] = True
for v in adj[u]:
if not seen[v]:
parent[v] = u
visit(v)
clock[0] += 1
f[u] = clock[0]
for u in range(n):
if not seen[u]:
visit(u)
return d, f, parent
def reachable(adj: list[list[int]], s: int) -> set[int]:
seen, todo = {s}, [s]
while todo:
u = todo.pop()
for v in adj[u]:
if v not in seen:
seen.add(v)
todo.append(v)
return seen
def test_example() -> None:
adj = [[1, 3, 2], [2], [0], [2, 4], [], [4, 6], [5]]
d, f, parent = depth_first_search(adj)
assert d == [1, 2, 3, 6, 7, 11, 12]
assert f == [10, 5, 4, 9, 8, 14, 13]
assert parent == [None, 0, 1, 0, 3, None, 5]
def test_edges() -> None:
assert depth_first_search([]) == ([], [], [])
assert depth_first_search([[0]]) == ([1], [2], [None])
assert depth_first_search([[], []]) == ([1, 3], [2, 4], [None, None])
def test_long_path() -> None:
n = 100000
adj = [[i + 1] for i in range(n - 1)] + [[]]
d, f, _ = depth_first_search(adj)
assert d[n - 1] == n and f[n - 1] == n + 1 and f[0] == 2 * n
def test_random_against_recursive() -> None:
sys.setrecursionlimit(10000)
rng = random.Random(7781)
for _ in range(600):
n = rng.randint(1, 12)
adj: list[list[int]] = [[] for _ in range(n)]
for _ in range(rng.randint(0, 3 * n)):
adj[rng.randrange(n)].append(rng.randrange(n))
d, f, parent = depth_first_search(adj)
assert (d, f, parent) == recursive_dfs(adj)
assert sorted(d + f) == list(range(1, 2 * n + 1))
for u in range(n):
for v in range(n):
nested = d[u] < d[v] < f[v] < f[u]
disjoint = f[u] < d[v] or f[v] < d[u]
assert u == v or nested or disjoint or d[v] < d[u] < f[u] < f[v]
for r in range(n):
if parent[r] is None:
tree = {v for v in range(n) if d[r] <= d[v] and f[v] <= f[r]}
assert tree <= reachable(adj, r)
What the next level needs
- the tests did not run in this build