Quick-union in Python
Union-find that links one root to another on union and follows parent pointers on find.
- File
quick_union.py, 65 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/quick-union/. The authors are listed in prov/TPV99M.jsonl.
class QuickUnion:
"""Union-find on the elements 0 to n - 1 as a forest of parent pointers.
find walks from an element to its root. union links the root of the
first argument under the root of the second. Both cost time in
proportion to the depth of the elements, which is O(n) in the worst case.
"""
def __init__(self, n: int) -> None:
if n < 0:
raise ValueError("n must be non-negative")
self.parent: list[int] = list(range(n)) # @line L1
self.count = n # @line L2
def __len__(self) -> int:
return len(self.parent)
def find(self, p: int) -> int:
"""Return the root of the tree of p."""
parent = self.parent
if not 0 <= p < len(parent):
raise IndexError("element out of range")
while p != parent[p]: # @line L3
p = parent[p] # @line L4
return p # @line L5
def union(self, p: int, q: int) -> bool:
"""Merge the sets of p and q. Return False if they were already one set."""
i, j = self.find(p), self.find(q) # @line L6
if i == j: # @line L7
return False
self.parent[i] = j # @line L8
self.count -= 1 # @line L9
return True
def connected(self, p: int, q: int) -> bool:
"""Return True if p and q are in the same set."""
return self.find(p) == self.find(q)
def depth(self, p: int) -> int:
"""Return the number of edges from p to its root."""
d = 0
while p != self.parent[p]:
p = self.parent[p]
d += 1
return d
def check_invariant(self) -> None:
"""Check that the parent pointers form a forest with count roots."""
n = len(self.parent)
roots = 0
for x in range(n):
assert 0 <= self.parent[x] < n
if self.parent[x] == x:
roots += 1
steps, y = 0, x
while y != self.parent[y]:
y = self.parent[y]
steps += 1
assert steps < n
assert roots == self.count
test_quick_union.py 65 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/quick-union/. The authors are listed in prov/TPV99M.jsonl.
import random
import pytest
from quick_union import QuickUnion
TINY = [(4, 3), (3, 8), (6, 5), (9, 4), (2, 1), (8, 9), (5, 0), (7, 2), (6, 1), (1, 0), (6, 7)]
def test_example() -> None:
uf = QuickUnion(10)
merged = []
for p, q in TINY:
merged.append(uf.union(p, q))
uf.check_invariant()
assert merged == [True] * 5 + [False] + [True] * 3 + [False, False]
assert uf.parent == [1, 1, 1, 8, 3, 0, 5, 1, 8, 8]
assert [uf.depth(x) for x in range(10)] == [1, 0, 1, 1, 2, 2, 3, 1, 0, 1]
assert uf.count == 2
assert uf.find(6) == 1 and uf.find(9) == 8
def test_edges() -> None:
uf = QuickUnion(0)
assert uf.count == 0 and len(uf) == 0
uf = QuickUnion(1)
assert uf.find(0) == 0
assert not uf.union(0, 0)
assert uf.count == 1
with pytest.raises(IndexError):
uf.find(1)
with pytest.raises(ValueError):
QuickUnion(-1)
def test_chain_worst_case() -> None:
n = 2000
uf = QuickUnion(n)
for k in range(1, n):
assert uf.depth(0) == k - 1
uf.union(0, k)
assert uf.depth(0) == n - 1
assert uf.find(0) == n - 1
uf.check_invariant()
def test_random_against_labels() -> None:
rng = random.Random(29)
for _ in range(600):
n = rng.randrange(1, 14)
uf = QuickUnion(n)
label = list(range(n))
for _ in range(rng.randrange(30)):
p, q = rng.randrange(n), rng.randrange(n)
if rng.random() < 0.5:
assert uf.union(p, q) == (label[p] != label[q])
old, new = label[p], label[q]
label = [new if v == old else v for v in label]
uf.check_invariant()
else:
assert uf.connected(p, q) == (label[p] == label[q])
assert uf.count == len(set(label))
What the next level needs
- the tests did not run in this build