Union-Find in Python: keep connected components without rescanning everything

Dynamic connectivity starts with a stream of links and a practical question: do these two items belong to the same component? New links arrive over time, and each link may join two groups or turn out to be redundant.

That is the dynamic connectivity problem. The usual input is a stream of pairs, such as (p, q), where each pair says that p is connected to q.^[1] The useful part is that connectivity is transitive: if 4 is linked to 7 and 7 is linked to 9, then 4 and 9 are in the same component even if no direct edge joins them.

The slow answer is easy to write

Start with one label per item. To join two components, scan the whole label array and change every occurrence of one label to the other. Checking a pair is cheap, but a successful join can touch every item.

That approach is perfectly respectable for a small input. Its cost becomes visible when the input grows because every successful merge revisits the whole label array.

The union-find idea

Union-find stores a forest of parent links. Every item eventually leads to a root, and two items are connected when their roots match. A find operation follows parent links to the root. A union operation joins two roots.

Two small choices keep the forest shallow. First, attach the smaller tree below the larger one. Second, while finding a root, make the visited items point directly to it. The first choice is weighted union; the second is path compression. Princeton’s algorithm notes describe the same progression from quick-union to weighted quick-union and then path compression.^[1] The reference implementation combines weighted union by size with full path compression.^[2]

Imagine the links (0, 1), (2, 3) and (1, 2). The first two pairs form two small trees. The third pair finds the root of 1 and the root of 2, then joins those roots. A later lookup for 0 can flatten the path it travels, so that path may cost less on a later lookup. The effect depends on which paths the workload actually visits.

Complexity on paper

The label-based version does constant work for connected, but a successful union scans n labels. Because there can be at most n - 1 successful unions, a sequence of m operations takes O(n min(m, n) + m) time in the worst case, which becomes O(n² + m) once the stream is long enough. Its storage is O(n).

Weighted union limits the depth of a tree to at most log n.^[1] With path compression added, the amortised cost of an intermixed sequence of m find and union operations is O(m α(n)), where α is the inverse Ackermann function.^[1] For ordinary program sizes, α(n) grows so slowly that the bound behaves like a small constant, although the constant still depends on the language and the implementation.

Those are asymptotic bounds, not a promise that every short Python program will be faster. Union-find allocates parent and size arrays, performs several Python-level operations per lookup and needs a suitable workload to repay its setup cost. The benchmark below keeps input generation outside the timed section and checks accepted pairs, component count and the final component partition for both implementations.

A runnable comparison

Save this as union_find_benchmark.py and run it with Python 3. It uses one fixed random seed, four input sizes and three timed samples per implementation. The reported time is the median sample. Memory is the peak traced by tracemalloc, so it covers Python allocations made during the run rather than the whole process resident set.^[4]

from __future__ import annotations
import platform
import random
import time
import tracemalloc
from statistics import median
SEED = 20260820
SIZES = (500, 1_000, 2_000, 5_000)
REPEATS = 3
def make_edges(n: int, factor: int = 4) -> list[tuple[int, int]]:
    rng = random.Random(SEED + n)
    return [(rng.randrange(n), rng.randrange(n)) for _ in range(factor * n)]
class NaiveConnectivity:
    """Relabel every member of a component after each successful union."""
    def __init__(self, n: int) -> None:
        self.label = list(range(n))
        self.components = n
    def connected(self, p: int, q: int) -> bool:
        return self.label[p] == self.label[q]
    def union(self, p: int, q: int) -> bool:
        left = self.label[p]
        right = self.label[q]
        if left == right:
            return False
        for i, label in enumerate(self.label):
            if label == right:
                self.label[i] = left
        self.components -= 1
        return True
def find(parent: list[int], p: int) -> int:
    root = p
    while root != parent[root]:
        root = parent[root]
    while p != root:
        parent_of_p = parent[p]
        parent[p] = root
        p = parent_of_p
    return root
class WeightedPathCompression:
    """Union by size, with path compression during find."""
    def __init__(self, n: int) -> None:
        self.parent = list(range(n))
        self.size = [1] * n
        self.components = n
    def connected(self, p: int, q: int) -> bool:
        return find(self.parent, p) == find(self.parent, q)
    def union(self, p: int, q: int) -> bool:
        root_p = find(self.parent, p)
        root_q = find(self.parent, 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 partition_signature(structure, n: int) -> tuple[int, ...]:
    """Return component membership with roots renamed by first appearance."""
    if isinstance(structure, NaiveConnectivity):
        roots = structure.label
    elif isinstance(structure, WeightedPathCompression):
        roots = [find(structure.parent, item) for item in range(n)]
    else:
        raise TypeError(f"unsupported implementation: {type(structure).__name__}")
    names: dict[int, int] = {}
    signature = []
    for root in roots:
        if root not in names:
            names[root] = len(names)
        signature.append(names[root])
    return tuple(signature)
def process(
    implementation,
    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 not structure.connected(p, q):
            structure.union(p, q)
            accepted += 1
    result = (accepted, structure.components)
    if verify:
        return (*result, partition_signature(structure, n))
    return result
def time_process(implementation, n: int, edges: list[tuple[int, int]]) -> list[float]:
    process(implementation, n, edges)  # warm-up run
    samples = []
    for _ in range(REPEATS):
        start = time.perf_counter()
        result = process(implementation, n, edges)
        samples.append(time.perf_counter() - start)
        if result != process(implementation, n, edges):
            raise AssertionError("non-deterministic result")
    return samples
def peak_memory(implementation, n: int, edges: list[tuple[int, int]]) -> int:
    tracemalloc.start()
    process(implementation, n, edges)
    _, peak = tracemalloc.get_traced_memory()
    tracemalloc.stop()
    return peak
def main() -> None:
    print(f"Python {platform.python_version()} on {platform.platform()}")
    print(f"seed={SEED}, edges=4*n, repeats={REPEATS}, timer=perf_counter")
    print("correctness_check=accepted_pairs,components,final_partition")
    print("n,naive_median_s,optimised_median_s,speedup,naive_peak_KiB,optimised_peak_KiB,accepted,components")
    for n in SIZES:
        edges = make_edges(n)
        baseline = process(NaiveConnectivity, n, edges, verify=True)
        optimised = process(WeightedPathCompression, n, edges, verify=True)
        if baseline != optimised:
            raise AssertionError(f"result mismatch for n={n}: {baseline} != {optimised}")
        naive_samples = time_process(NaiveConnectivity, n, edges)
        optimised_samples = time_process(WeightedPathCompression, n, edges)
        naive_median = median(naive_samples)
        optimised_median = median(optimised_samples)
        speedup = naive_median / optimised_median
        naive_peak = peak_memory(NaiveConnectivity, n, edges)
        optimised_peak = peak_memory(WeightedPathCompression, n, edges)
        print(
            f"{n},{naive_median:.6f},{optimised_median:.6f},{speedup:.2f},"
            f"{naive_peak / 1024:.1f},{optimised_peak / 1024:.1f},"
            f"{baseline[0]},{baseline[1]}"
        )
if __name__ == "__main__":
    main()

The code uses perf_counter, Python’s performance counter for measuring short durations.^[3] It does not claim to model a server, a different Python build or a different mix of connections. It measures this workload on the machine that ran it.

What the run showed

The result was:

Python 3.13.5 on Linux-6.12.47+rpt-rpi-v8-aarch64-with-glibc2.41
seed=20260820, edges=4*n, repeats=3, timer=perf_counter
correctness_check=accepted_pairs,components,final_partition
n,naive_median_s,optimised_median_s,speedup,naive_peak_KiB,optimised_peak_KiB,accepted,components
500,0.028519,0.002897,9.85,12.1,15.8,499,1
1000,0.118883,0.006040,19.68,31.5,39.1,999,1
2000,0.503199,0.012105,41.57,70.5,85.9,1997,3
5000,3.507390,0.030677,114.33,187.6,226.5,4994,6

On this run, the optimised implementation was about 9 times faster at 500 items and about 111 times faster at 5,000 items. The widening gap is consistent with the relabelling version repeatedly scanning a larger array while weighted union and path compression keep the parent forest shallow. This benchmark did not instrument scan counts or tree heights, so that explanation is based on the algorithms and the observed timings rather than a separate causal measurement. The optimised version used more traced memory, roughly 226 KiB versus 188 KiB at 5,000 items, because it keeps both parent and size arrays.

The numbers are ratios from one Linux aarch64 machine running Python 3.13.5, not universal conversion factors. The input has 4*n pseudorandom pairs, generated with seed 20260820. Both methods produced the same accepted-pair count, component count and final partition for every size. That is a useful check for this input, not a proof that either implementation handles every possible input.

When to use it

Use union-find for undirected incremental connectivity when connections are added and you need many connectivity checks. It fits incremental network construction, maze connectivity and connected-component labelling particularly well. Directed reachability, weighted constraints and frequent edge updates require different algorithms or data structures.

It is a poor fit when links must be removed frequently. Deleting an edge can split a component, and this simple structure has no cheap way to rebuild that history. The standard dynamic-connectivity discussion treats fully dynamic graphs as a more complicated problem than the incremental case.^[1]

For a handful of items, the label scan may be the better engineering choice because it is shorter and easier to inspect. For a large stream of merges and queries, the extra arrays buy a large reduction in repeated work. Measure with your own data before changing a hot path. Algorithms provide the shape of the cost; your workload supplies the actual numbers.