Aldous-Broder algorithm in Python
Builds a uniform random spanning tree from the first-visit edges of a random walk.
- File
aldous_broder.py, 93 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/aldous-broder/. The authors are listed in prov/KAFQ9E.jsonl.
import random
def aldous_broder(
n: int,
edges: list[tuple[int, int, float]],
rng: random.Random,
weighted: bool = False,
) -> list[int]:
"""Return a random spanning forest as a sorted list of edge indices.
The vertices are 0 .. n - 1 and edges[i] = (u, v, w) is an undirected
edge. With weighted = False the weights are ignored and the forest is
uniform among the spanning forests with the components of the graph
(a uniform spanning tree when the graph is connected). With
weighted = True each w must be positive, and a forest F has probability
proportional to the product of the weights of its edges.
Parallel edges are different edges. Self-loops are allowed; the walk may
take them, but they never enter the forest. Each component is handled by
its own walk, started at its smallest vertex.
Expected cost: the cover time of the walk, at most 2 m (n - 1) steps for
a connected unweighted graph; O(n + m) space.
"""
adj: list[list[int]] = [[] for _ in range(n)]
for i, (u, v, _) in enumerate(edges):
adj[u].append(i)
if v != u:
adj[v].append(i)
# Cumulative weights per vertex, so that a step picks an edge in O(log deg).
cum: list[list[float]] = []
for x in range(n):
acc, row = 0.0, []
for i in adj[x]:
u, v, w = edges[i]
c = (w if weighted else 1.0) * (2 if u == v else 1) # a loop has two ends at x
acc += c
row.append(acc)
cum.append(row)
if weighted and any(w <= 0 for _, _, w in edges):
raise ValueError("weights must be positive when weighted is True")
seen = [False] * n
tree: list[int] = []
for start in range(n):
if seen[start]:
continue
size = _component_size(adj, edges, start)
x = start # @line L1
seen[x] = True # @line L2
visited = 1
while visited < size: # @line L3
i = _step(adj[x], cum[x], rng) # @line L4
u, v, _ = edges[i]
y = v if u == x else u
if not seen[y]: # @line L5
seen[y] = True # @line L6
tree.append(i) # @line L7
visited += 1
x = y # @line L8
return sorted(tree) # @line L9
def _step(inc: list[int], cum: list[float], rng: random.Random) -> int:
"""Pick an incident edge with probability proportional to its weight."""
r = rng.random() * cum[-1]
lo, hi = 0, len(cum) - 1
while lo < hi:
mid = (lo + hi) // 2
if cum[mid] > r:
hi = mid
else:
lo = mid + 1
return inc[lo]
def _component_size(adj: list[list[int]], edges: list[tuple[int, int, float]], s: int) -> int:
"""Number of vertices reachable from s."""
mark = {s}
stack = [s]
while stack:
x = stack.pop()
for i in adj[x]:
u, v, _ = edges[i]
y = v if u == x else u
if y not in mark:
mark.add(y)
stack.append(y)
return len(mark)
test_aldous_broder.py 104 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/aldous-broder/. The authors are listed in prov/KAFQ9E.jsonl.
import itertools
import math
import random
from collections import Counter
from aldous_broder import aldous_broder
EXAMPLE = [(0, 1, 4), (0, 2, 3), (1, 2, 1), (1, 3, 2), (2, 3, 4), (2, 4, 5), (3, 4, 7), (3, 5, 6), (4, 5, 2)]
def components(n, edges, chosen):
parent = list(range(n))
def find(x):
while parent[x] != x:
x = parent[x]
return x
for i in chosen:
u, v, _ = edges[i]
parent[find(u)] = find(v)
return len({find(x) for x in range(n)})
def spanning_forests(n, edges):
"""All edge sets that are forests with the components of the graph."""
c = components(n, edges, range(len(edges)))
k = n - c
return [
s
for s in itertools.combinations(range(len(edges)), k)
if components(n, edges, s) == c
]
def chi_square(counts, expected):
return sum((counts.get(t, 0) - e) ** 2 / e for t, e in expected.items())
def check_distribution(n, edges, samples, seed, weighted=False):
forests = spanning_forests(n, edges)
weight = {s: (math.prod(edges[i][2] for i in s) if weighted else 1.0) for s in forests}
total = sum(weight.values())
expected = {s: samples * weight[s] / total for s in forests}
rng = random.Random(seed)
counts = Counter(tuple(aldous_broder(n, edges, rng, weighted)) for _ in range(samples))
assert set(counts) <= set(forests)
df = len(forests) - 1
# A bound far in the upper tail: mean df, standard deviation sqrt(2 df).
assert chi_square(counts, expected) < df + 6 * math.sqrt(2 * df) + 10
def test_uniform_on_example():
assert len(spanning_forests(6, EXAMPLE)) == 55
check_distribution(6, EXAMPLE, 55 * 200, seed=1990)
def test_uniform_on_k4_and_with_parallel_edges_and_loops():
k4 = [(i, j, 1) for i in range(4) for j in range(i + 1, 4)]
check_distribution(4, k4, 16 * 300, seed=1)
multi = [(0, 1, 1), (0, 1, 1), (1, 2, 1), (2, 0, 1), (2, 2, 1)]
check_distribution(3, multi, 5 * 600, seed=2)
def test_weighted_distribution():
tri = [(0, 1, 1.0), (1, 2, 2.0), (2, 0, 3.0)]
check_distribution(3, tri, 6000, seed=3, weighted=True)
check_distribution(6, EXAMPLE, 30000, seed=4, weighted=True)
def test_disconnected_gives_uniform_forest():
g = [(0, 1, 1), (1, 2, 1), (2, 0, 1), (3, 4, 1), (4, 5, 1), (5, 3, 1), (3, 4, 1)]
check_distribution(7, g, 15 * 300, seed=5)
def test_edge_cases():
rng = random.Random(0)
assert aldous_broder(0, [], rng) == []
assert aldous_broder(1, [], rng) == []
assert aldous_broder(1, [(0, 0, 1)], rng) == []
assert aldous_broder(2, [], rng) == []
assert aldous_broder(2, [(0, 1, -3)], rng) == [0]
try:
aldous_broder(2, [(0, 1, 0)], rng, weighted=True)
except ValueError:
pass
else:
raise AssertionError("expected ValueError")
def test_random_graphs_give_spanning_forests():
rng = random.Random(1989)
for _ in range(600):
n = rng.randrange(1, 8)
m = rng.randrange(0, 12)
edges = [(rng.randrange(n), rng.randrange(n), rng.choice([-1, 0, 2, 2])) for _ in range(m)]
f = aldous_broder(n, edges, rng)
c = components(n, edges, range(m))
assert len(f) == n - c
assert components(n, edges, f) == c
assert all(edges[i][0] != edges[i][1] for i in f)
What the next level needs
- the tests did not run in this build