Union-find in Python: when relabelling every component stops scaling

Problem. A dynamic-connectivity stream supplies pairs of integer IDs and asks whether a new pair joins two components or repeats a connection.^[1] The implementation must keep the component state correct while the stream is still arriving.

Baseline. QuickFind.union() stores one label per item. When two labels differ, the shown loop scans all n labels and rewrites every entry belonging to one component.^[1] A successful merge therefore does O(n) work in this exact Python implementation.

Solution. Weighted parent trees keep each item on a path to a root, while path compression rewrites the links visited by find(). The smaller tree is attached below the larger tree, following the structure of the reference weighted quick-union implementation.^[1][2]

Measured result. In the retained run, 20,000 generated pairs over 5,000 items took a median of 3.235069 seconds with quick-find and 0.028589 seconds with weighted union plus path compression. The ratio was 113.16x. Both implementations produced the same complete component partition; their traced peaks were 187.6 KiB and 226.5 KiB respectively.

Reproduce the result

Complete code

from __future__ import annotations
import hashlib
import platform
import random
import time
import tracemalloc
from statistics import median
SEED = 20261007
EDGE_FACTOR = 4
SIZES = (500, 1_000, 2_000, 5_000)
REPEATS = 5
class QuickFind:
    """Baseline: relabel every item in the right component after a union."""
    def __init__(self, n: int) -> None:
        self.label = list(range(n))
        self.components = n
    def union(self, p: int, q: int) -> bool:
        left = self.label[p]
        right = self.label[q]
        if left == right:
            return False
        for index, label in enumerate(self.label):
            if label == right:
                self.label[index] = left
        self.components -= 1
        return True
    def roots(self) -> list[int]:
        return self.label
class WeightedPathCompression:
    """Solution: weighted parent trees plus full path compression."""
    def __init__(self, n: int) -> None:
        self.parent = list(range(n))
        self.size = [1] * n
        self.components = n
    def find(self, item: int) -> int:
        root = item
        while root != self.parent[root]:
            root = self.parent[root]
        while item != root:
            parent_of_item = self.parent[item]
            self.parent[item] = root
            item = parent_of_item
        return root
    def union(self, p: int, q: int) -> bool:
        root_p = self.find(p)
        root_q = self.find(q)
        if root_p == root_q:
            return False
        if self.size[root_p] < self.size[root_q]:
            root_p, root_q = root_q, root_p
        self.parent[root_q] = root_p
        self.size[root_p] += self.size[root_q]
        self.components -= 1
        return True
    def roots(self) -> list[int]:
        return [self.find(item) for item in range(len(self.parent))]
def make_edges(n: int) -> list[tuple[int, int]]:
    rng = random.Random(SEED + n)
    return [(rng.randrange(n), rng.randrange(n)) for _ in range(EDGE_FACTOR * n)]
def partition_signature(structure: QuickFind | WeightedPathCompression) -> tuple[int, ...]:
    names: dict[int, int] = {}
    signature: list[int] = []
    for root in structure.roots():
        if root not in names:
            names[root] = len(names)
        signature.append(names[root])
    return tuple(signature)
def process(
    implementation: type[QuickFind] | type[WeightedPathCompression],
    n: int,
    edges: list[tuple[int, int]],
    *,
    verify: bool = False,
) -> tuple[int, int] | tuple[int, int, tuple[int, ...]]:
    structure = implementation(n)
    accepted = 0
    for p, q in edges:
        if structure.union(p, q):
            accepted += 1
    result = (accepted, structure.components)
    if verify:
        return (*result, partition_signature(structure))
    return result
def timed_samples(
    implementation: type[QuickFind] | type[WeightedPathCompression],
    n: int,
    edges: list[tuple[int, int]],
) -> list[float]:
    expected = process(implementation, n, edges)
    process(implementation, n, edges)  # one untimed warm-up
    samples: list[float] = []
    for _ in range(REPEATS):
        start = time.perf_counter()
        result = process(implementation, n, edges)
        samples.append(time.perf_counter() - start)
        if result != expected:
            raise AssertionError("non-deterministic result")
    return samples
def peak_traced_memory(
    implementation: type[QuickFind] | type[WeightedPathCompression],
    n: int,
    edges: list[tuple[int, int]],
) -> int:
    tracemalloc.start()
    process(implementation, n, edges)
    _, peak = tracemalloc.get_traced_memory()
    tracemalloc.stop()
    return peak
def digest(signature: tuple[int, ...]) -> str:
    return hashlib.sha256(repr(signature).encode("ascii")).hexdigest()[:16]
def trace_example() -> None:
    structure = WeightedPathCompression(4)
    for p, q in ((0, 1), (2, 3), (0, 2)):
        structure.union(p, q)
    before = structure.parent.copy()
    root = structure.find(3)
    after = structure.parent.copy()
    print(f"trace_parent_before={before}")
    print(f"trace_find_3={root}")
    print(f"trace_parent_after={after}")
def edge_case_checks() -> None:
    for implementation in (QuickFind, WeightedPathCompression):
        empty = implementation(0)
        assert empty.components == 0
        assert partition_signature(empty) == ()
        small = implementation(3)
        assert not small.union(1, 1)
        assert small.union(0, 1)
        assert small.union(1, 2)
        assert small.components == 1
        assert partition_signature(small) == (0, 0, 0)
def main() -> None:
    print(f"Python {platform.python_version()} on {platform.platform()}")
    print(
        f"seed={SEED}, edges={EDGE_FACTOR}*n, repeats={REPEATS}, "
        "timer=time.perf_counter"
    )
    print("correctness=accepted_pairs,components,full_partition")
    trace_example()
    print("n,quick_find_median_s,weighted_pc_median_s,speedup,quick_find_peak_KiB,weighted_pc_peak_KiB")
    for n in SIZES:
        edges = make_edges(n)
        quick_find = process(QuickFind, n, edges, verify=True)
        weighted_pc = process(WeightedPathCompression, n, edges, verify=True)
        if quick_find != weighted_pc:
            raise AssertionError(f"result mismatch for n={n}")
        quick_samples = timed_samples(QuickFind, n, edges)
        weighted_samples = timed_samples(WeightedPathCompression, n, edges)
        quick_median = median(quick_samples)
        weighted_median = median(weighted_samples)
        speedup = quick_median / weighted_median
        quick_peak = peak_traced_memory(QuickFind, n, edges)
        weighted_peak = peak_traced_memory(WeightedPathCompression, n, edges)
        accepted, components, signature = quick_find
        print(f"n={n}")
        print("quick_find_samples=" + repr([round(value, 6) for value in quick_samples]))
        print("weighted_pc_samples=" + repr([round(value, 6) for value in weighted_samples]))
        print(f"quick_find_min={min(quick_samples):.6f}")
        print(f"weighted_pc_min={min(weighted_samples):.6f}")
        print(f"quick_find_median={quick_median:.6f}")
        print(f"weighted_pc_median={weighted_median:.6f}")
        print(f"speedup={speedup:.2f}")
        print(
            f"result=accepted:{accepted},components:{components},"
            f"partition_digest:{digest(signature)}"
        )
        print(
            f"peak_KiB=quick_find:{quick_peak / 1024:.1f},"
            f"weighted_pc:{weighted_peak / 1024:.1f}"
        )
    edge_case_checks()
    print("edge_checks=pass")
if __name__ == "__main__":
    main()

Command

python3 union_find_benchmark.py

Output

$ python3 union_find_benchmark.py
Python 3.13.5 on Linux-6.12.47+rpt-rpi-v8-aarch64-with-glibc2.41
seed=20261007, edges=4*n, repeats=5, timer=time.perf_counter
correctness=accepted_pairs,components,full_partition
trace_parent_before=[0, 0, 0, 2]
trace_find_3=0
trace_parent_after=[0, 0, 0, 0]
n,quick_find_median_s,weighted_pc_median_s,speedup,quick_find_peak_KiB,weighted_pc_peak_KiB
n=500
quick_find_samples=[0.026659, 0.026618, 0.029295, 0.027139, 0.026526]
weighted_pc_samples=[0.002835, 0.002902, 0.002683, 0.003176, 0.003366]
quick_find_min=0.026526
weighted_pc_min=0.002683
quick_find_median=0.026659
weighted_pc_median=0.002902
speedup=9.19
result=accepted:499,components:1,partition_digest:7da468df52104ca8
peak_KiB=quick_find:12.1,weighted_pc:15.8
n=1000
quick_find_samples=[0.116092, 0.1191, 0.116106, 0.119232, 0.116171]
weighted_pc_samples=[0.005666, 0.005532, 0.005533, 0.005797, 0.006273]
quick_find_min=0.116092
weighted_pc_min=0.005532
quick_find_median=0.116171
weighted_pc_median=0.005666
speedup=20.50
result=accepted:999,components:1,partition_digest:142d103e42368fc5
peak_KiB=quick_find:31.5,weighted_pc:39.1
n=2000
quick_find_samples=[0.498821, 0.49042, 0.49238, 0.518675, 0.505865]
weighted_pc_samples=[0.011483, 0.011503, 0.011286, 0.011519, 0.011562]
quick_find_min=0.490420
weighted_pc_min=0.011286
quick_find_median=0.498821
weighted_pc_median=0.011503
speedup=43.37
result=accepted:1998,components:2,partition_digest:39211bf5feef0fef
peak_KiB=quick_find:70.5,weighted_pc:85.9
n=5000
quick_find_samples=[3.361603, 3.179151, 3.235069, 3.239637, 3.16423]
weighted_pc_samples=[0.027878, 0.027981, 0.029971, 0.033744, 0.028589]
quick_find_min=3.164230
weighted_pc_min=0.027878
quick_find_median=3.235069
weighted_pc_median=0.028589
speedup=113.16
result=accepted:4995,components:5,partition_digest:7e7dcc8a3b3f547c
peak_KiB=quick_find:187.6,weighted_pc:226.5
edge_checks=pass
exit=0

Environment

The run used CPython 3.13.5 on Linux 6.12.47, aarch64, with only the Python standard library. Input generation happened before timing. The source and captured transcript are retained with the experiment. The timing vectors in the captured output are rounded to six decimal places; medians, minima and speedups are calculated from the full-precision samples before display.

Methodology

make_edges() creates 4 * n pairs with a private random.Random instance seeded from 20261007 + n. The tested sizes are 500, 1,000, 2,000 and 5,000.

Python documents Random instances as independent generators, and an integer seed makes this input construction repeatable for the same Python implementation.^[5]

Each implementation receives the same retained edge list. The correctness pass records accepted unions, remaining component count and a full partition signature, then compares all three values before timing. The printed partition digest is a compact audit value; the comparison itself uses the complete tuple.

Each size gets one untimed warm-up, then five timed samples. The timed section creates the structure and processes all pairs. It excludes input generation, partition verification and memory tracing.

The timer is time.perf_counter(), which Python documents as a performance counter for short-duration measurements.^[3]

The memory pass starts tracemalloc immediately before a fresh run and records its peak after the run. These figures are traced Python allocations for this call, not resident-set size, allocator-wide memory or native allocations.

tracemalloc traces Python memory blocks, according to the Python documentation.^[4]

Limits

This is one run on one Linux aarch64 host. The timings show the direction and size of the difference for this input generator and operation contract. They do not provide a portable speed factor for another CPU, interpreter, background load or object type.

The edge stream is pseudo-random rather than a retained production trace. It contains integer pairs, including repeated and self pairs, and the code assumes every endpoint is in range(n). The explicit edge checks cover an empty structure and repeated or merging pairs, not invalid endpoint handling.

The comparison measures quick-find against weighted union with full path compression. It does not measure quick-union without weighting, path halving, a third-party library, edge deletion, directed reachability or weighted constraints.

The dynamic-connectivity problem here is incremental: pairs arrive and components merge.^[1]

The baseline loop and the optimised structure both use Python lists. The optimised version keeps a parent list and a size list, so its extra traced memory is expected for this representation. tracemalloc starts after imports and input construction, and its peak must not be read as the process’s full memory footprint.

The problem

The input model is deliberately small: an item is an integer ID, and a pair (p, q) says that the two items should belong to the same equivalence class.^[1]

Once p is connected to q, and q to r, a later query must treat p and r as connected even without a direct pair.^[1]

The benchmark turns each input pair into one union() call. A return value of True means that the pair joined two previously separate components. A return value of False means that the pair was already implied by the current partition. This keeps the outer workload identical while the two classes maintain the state differently.

Baseline: quick-find relabels the array

Quick-find makes a connectivity check cheap: two items are in the same component when their labels match.^[1] The cost moves into union(). When left and right differ, the baseline scans every slot and changes each right label to left.

For m input pairs, no more than n - 1 of them can be successful merges. The exact loop therefore has worst-case relabelling work of O(n min(m, n) + m), with O(n) storage. The bound is more informative than simply writing O(nm), because successful merges stop once one component contains every item.

The array is easy to inspect, which is useful in a short teaching example. It also does the same full scan when a component has only one member left to relabel. That repeated work is the part the second representation removes.

Solution: weighted trees and path compression

The second class stores a parent link for each item. A root points to itself. To find an item’s component, find() follows parent links until it reaches a root. Two items share a component when those roots match.^[1]

The reference implementation uses weighted quick union by size with full path compression.^[2]

The size list records the number of items in each root’s tree. On a successful union, the smaller root points to the larger root. This limits how tall a tree can become under weighted union.^[1]

Path compression runs during find(). After the root is known, the second loop sends every visited item directly to that root. The small trace in the output starts with 3 -> 2 -> 0; after find(3), the parent for 3 is 0. The operation changes the path it used, rather than scanning unrelated items.

What the measurement shows

The gap widens across the tested sizes. At n=500, the quick-find median was 0.026659 seconds and the weighted path-compression median was 0.002902 seconds, a 9.19x ratio. At n=1,000, the medians were 0.116171 and 0.005666 seconds, a 20.50x ratio.

At n=2,000, the medians were 0.498821 and 0.011503 seconds, a 43.37x ratio. At n=5,000, they were 3.235069 and 0.028589 seconds, a 113.16x ratio. The raw output keeps all five samples and the minimum for each implementation, so the medians are not the only timing evidence retained locally.

The complete partition check passed at every size. At n=5,000, both versions accepted 4,995 pairs and ended with five components. Those counters alone would not prove that every item had the same component membership, so the code also compared the full normalised partition signature. The digest printed in the transcript is only a compact record of that signature.

The weighted representation used more traced memory at every reported size. At n=5,000 its peak was 226.5 KiB, compared with 187.6 KiB for quick-find. That is the measured cost of this implementation’s two parent-state lists under the stated tracing boundary, not a general memory ratio for union-find.

The timing pattern is consistent with the two algorithms’ work: successful quick-find merges keep scanning the whole label array, while the weighted implementation follows and repairs parent paths. The run did not count label writes, parent traversals or tree heights, so that sentence is an algorithmic interpretation of the measured result, not a separate causal instrument reading.

Practical applications

– Incremental network or asset connectivity: map each device, service or asset to an integer ID and call union(p, q) as links become known. The same component count and root checks can answer whether two IDs have joined. The benchmark measures the merge workload, not network discovery, input parsing or concurrent access.

– Maze and grid-component labelling: assign one ID to each cell and union adjacent open cells while scanning the grid. The parent array represents the component relation without rescanning every labelled cell after each join. The article does not measure grid construction or image-processing throughput.

– Kruskal-style cycle checks: after sorting candidate edges by weight, use the root of each endpoint to decide whether adding an edge would join two components. Sorting and edge ordering dominate parts of that workload, so the timings here should not be presented as a measurement of a minimum-spanning-tree implementation.

– Duplicate relation filtering: a stream of pairs can be reduced to the pairs that actually merge separate components, matching the accepted-pair counter in the code. This is useful when later stages need only component-joining events, though removal of an earlier relation is outside this data structure’s contract.

For this exact 20,000-pair run, the optimised object ended with five components and a 113.16x median ratio at n=5,000. A different stream can change both the accepted-pair pattern and the amount of path work, so the retained benchmark is a starting point for measuring a real input rather than a replacement for that measurement.

Sources

[1] Princeton Algorithms 4th Edition: Case Study, Union-Find

[2] Princeton Algorithms 4th Edition: WeightedQuickUnionPathCompressionUF.java

[3] Python 3.13 documentation: time — Time access and conversions

[4] Python 3.13 documentation: tracemalloc — Trace memory allocations

[5] Python 3.13 documentation: random — Generate pseudo-random numbers