Wilson's algorithm in Python
Builds a uniform random spanning tree from loop-erased random walks.
- File
wilson_algorithm.py, 102 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/wilson-algorithm/. The authors are listed in prov/XQWGT5.jsonl.
import random
def wilson(
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; a walk may
take them, but loop erasure removes them. The root of each component is
its smallest vertex, and the walks start from the other vertices in
increasing order.
Expected cost: O(m n) random steps for a connected graph (the commute
time between the root and a vertex is at most 2 m (n - 1)); O(n + m)
space.
"""
if weighted and any(w <= 0 for _, _, w in edges):
raise ValueError("weights must be positive when weighted is True")
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)
in_tree = [False] * n
nxt = [-1] * n # the edge of the last step taken out of each vertex
tree: list[int] = []
comp = _components(n, adj, edges)
for x in range(n):
if comp[x] == x:
in_tree[x] = True # @line L1
for i in range(n): # @line L2
u = i # @line L3
while not in_tree[u]: # @line L4
nxt[u] = _step(adj[u], cum[u], rng) # @line L5
u = _other(edges[nxt[u]], u) # @line L6
u = i # @line L7
while not in_tree[u]: # @line L8
in_tree[u] = True # @line L9
tree.append(nxt[u]) # @line L10
u = _other(edges[nxt[u]], u) # @line L11
return sorted(tree) # @line L12
def _other(edge: tuple[int, int, float], x: int) -> int:
u, v, _ = edge
return v if u == x else u
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 _components(n: int, adj: list[list[int]], edges: list[tuple[int, int, float]]) -> list[int]:
"""The smallest vertex of the component of each vertex."""
comp = [-1] * n
for s in range(n):
if comp[s] != -1:
continue
comp[s] = s
stack = [s]
while stack:
x = stack.pop()
for i in adj[x]:
y = _other(edges[i], x)
if comp[y] == -1:
comp[y] = s
stack.append(y)
return comp
test_wilson_algorithm.py 129 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/wilson-algorithm/. The authors are listed in prov/XQWGT5.jsonl.
import itertools
import math
import random
from collections import Counter
from wilson_algorithm import wilson
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)]
class Scripted:
"""A stand-in for random.Random that returns fixed values."""
def __init__(self, values):
self.values = list(values)
def random(self):
return self.values.pop(0)
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 (brute force)."""
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(wilson(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_example_trace():
# b -> c -> b (loop erased) -> a; c -> b; d -> c; e -> f -> e (loop erased) -> d; f -> e.
rng = Scripted([0.5, 0.375, 0.1, 0.375, 0.375, 0.9, 0.75, 0.5, 0.75])
assert wilson(6, EXAMPLE, rng) == [0, 2, 4, 6, 8]
assert rng.values == []
def test_uniform_on_example():
assert len(spanning_forests(6, EXAMPLE)) == 55
check_distribution(6, EXAMPLE, 55 * 200, seed=1996)
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_ties_in_weights():
# Equal weights give equal factors, so this is the uniform law.
c4 = [(0, 1, 2.0), (1, 2, 2.0), (2, 3, 2.0), (3, 0, 2.0), (0, 2, 2.0)]
check_distribution(4, c4, 8 * 400, seed=6, 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 wilson(0, [], rng) == []
assert wilson(1, [], rng) == []
assert wilson(1, [(0, 0, 1)], rng) == []
assert wilson(2, [], rng) == []
assert wilson(2, [(0, 1, -3)], rng) == [0]
assert wilson(3, [(0, 1, 1), (1, 2, 1)], rng) == [0, 1]
try:
wilson(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(1996)
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 = wilson(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)
assert f in [list(s) for s in spanning_forests(n, edges)]
What the next level needs
- the tests did not run in this build