Union-Find (Disjoint Set Union)

Intuition

Union-find (also called a disjoint set union, DSU) keeps track of how elements are partitioned into groups while groups keep merging. It answers two questions fast:

• find(x): which group is x in? The answer is a representative element, the root of x's tree.
• union(a, b): merge the groups containing a and b.

Two elements are in the same group exactly when find(a) == find(b).

Internally each group is a tree stored in a parent array; a root is its own parent. find walks up parent links to the root, and union links one root under the other. Left alone, those trees can degenerate into long chains and find becomes O(n). Two small tricks fix that:

• Path compression: while walking up during find, re-point nodes closer to the root, so the next lookup is almost direct. The full version points every visited node at the root; path halving (point each node to its grandparent) is a one-pass iterative variant with the same guarantees.
• Union by size (or rank): always hang the smaller tree under the root of the larger one. A node only gets deeper when its group at least doubles in size, so tree height stays at most log2(n) even without compression.

With both tricks, any sequence of m operations on n elements runs in O(m * α(n)), where α is the inverse Ackermann function. It grows so slowly that α(n) ≤ 4 for any input that fits in the universe, so treat each operation as constant time. Union by size and union by rank give the same bound (log2 of the size behaves like the rank). Path compression alone is O(log n) amortized in the worst case, and union by size alone is O(log n) per operation; both are usually fine in practice, but write both tricks in interviews, since it costs two extra lines.

One limitation matters a lot: union-find cannot split a group efficiently. If edges are deleted, look for an offline trick (process the operations in reverse so deletions become additions).

When to use it in interviews / OAs

• Connectivity that only grows: "are a and b connected?", number of connected components or provinces, number of islands after each land addition, accounts merge (link records that share an email).
• Cycle detection in an undirected graph: an edge whose endpoints are already connected closes a cycle (redundant connection, "is this graph a valid tree?").
• Kruskal's minimum spanning tree: sort edges by weight, keep an edge if union succeeds.
• Equivalence classes: equality equations (apply all == first, then check every !=), similar strings, synonyms, "swap any characters at linked indices" (sort within each component).
• Offline threshold queries: sort edges and queries by weight, add edges up to each query's limit, then answer connectivity.
• Relations with weights: ratios between variables, parity or difference constraints; use a weighted DSU.

If the graph is static and you just need components once, BFS/DFS works too. Union-find shines when edges arrive over time or when it keeps the code short.

Template

A complete DSU with union by size, path halving, a component counter and set sizes. The union method returns whether a merge happened, which is exactly what cycle detection and Kruskal need.

class DSU:
    def __init__(self, n: int) -> None:
        self.parent = list(range(n))  # Every element starts as its own root.
        self.size = [1] * n
        self.components = n

    def find(self, x: int) -> int:
        while self.parent[x] != x:
            self.parent[x] = self.parent[self.parent[x]]  # Path halving: skip to the grandparent.
            x = self.parent[x]
        return x

    def union(self, a: int, b: int) -> bool:
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False  # Already in the same set.
        if self.size[ra] < self.size[rb]:
            ra, rb = rb, ra
        self.parent[rb] = ra  # Union by size: hang the smaller tree under the larger root.
        self.size[ra] += self.size[rb]
        self.components -= 1
        return True

    def connected(self, a: int, b: int) -> bool:
        return self.find(a) == self.find(b)

    def set_size(self, x: int) -> int:
        return self.size[self.find(x)]


def redundant_edge(n: int, edges: list[tuple[int, int]]) -> tuple[int, int] | None:
    # First edge whose endpoints are already connected: it closes a cycle.
    dsu = DSU(n)
    for u, v in edges:
        if not dsu.union(u, v):
            return (u, v)
    return None


d = DSU(5)
d.union(0, 1)
d.union(3, 4)
assert d.connected(0, 1) and not d.connected(1, 3) and d.components == 3 and d.set_size(4) == 2
assert redundant_edge(4, [(0, 1), (1, 2), (2, 0), (2, 3)]) == (2, 0)

For problems with string or coordinate keys, either map them to integer ids first (a dict from key to index) or replace the lists with dicts. For an R x C grid, the cell (r, c) becomes r * C + c.

A weighted DSU additionally stores, for each node, its offset relative to its parent. Path compression folds the offsets along the path into a single offset to the root. It answers questions like "given a - b = 5 and b - c = 3, what is a - c?" and detects contradictions.

class WeightedDSU:
    # Tracks values known only through differences: value[x] - value[y] = d.
    def __init__(self, n: int) -> None:
        self.parent = list(range(n))
        self.size = [1] * n
        self.w = [0] * n  # w[x] = value[x] - value[parent[x]]

    def find(self, x: int) -> int:
        p = self.parent[x]
        if p == x:
            return x
        root = self.find(p)  # Recursion depth is O(log n) thanks to union by size.
        self.w[x] += self.w[p]  # w[p] is now relative to the root.
        self.parent[x] = root
        return root

    def union(self, x: int, y: int, d: int) -> bool:
        # Record value[x] - value[y] = d; return False if it contradicts earlier facts.
        rx, ry = self.find(x), self.find(y)
        if rx == ry:
            return self.w[x] - self.w[y] == d
        if self.size[rx] < self.size[ry]:
            rx, ry, x, y, d = ry, rx, y, x, -d
        self.parent[ry] = rx
        self.w[ry] = self.w[x] - self.w[y] - d  # Chosen so that value[x] - value[y] == d holds.
        self.size[rx] += self.size[ry]
        return True

    def diff(self, x: int, y: int) -> int | None:
        if self.find(x) != self.find(y):
            return None  # No chain of facts links x and y.
        return self.w[x] - self.w[y]


wd = WeightedDSU(4)
assert wd.union(0, 1, 5) and wd.union(1, 2, 3)  # v0 - v1 = 5, v1 - v2 = 3
assert wd.diff(0, 2) == 8 and wd.diff(2, 0) == -8 and wd.diff(0, 3) is None
assert not wd.union(0, 2, 7) and wd.union(2, 0, -8)

The same structure handles ratios (multiply instead of add: "a / b = 2.0, b / c = 3.0, what is a / c?") and parity (add modulo 2: "x and y are in different groups", which checks bipartiteness online).

Complexity

• find / union: O(α(n)) amortized with both optimizations, effectively O(1).
• Building the structure: O(n). Space: O(n).
• Kruskal: O(E log E), dominated by sorting the edges.
• Recursive find is safe in Python only when union by size/rank keeps the depth at O(log n); otherwise use the iterative version.

Common pitfalls

• Comparing parents instead of roots. parent[a] == parent[b] is wrong; always compare find(a) == find(b).
• Linking non-roots. parent[b] = a without calling find first can corrupt the forest. Link find(b) under find(a).
• Counting components wrongly. Decrement the counter only when a union actually merges two different groups.
• Deep recursion. A recursive find without union by size can hit Python's recursion limit on a long chain.
• Index mapping bugs. Mixing 0- and 1-based ids, or forgetting to allocate an id for a new key, gives silent wrong answers.
• Expecting deletions to work. Removing an edge does not un-merge anything. Reverse time, or use a different structure.
• Weighted DSU sign errors. Fix the meaning of the stored offset (here value[x] - value[parent[x]]) and derive the root link formula from it; do not guess the signs.

Interview variations

• Number of provinces / connected components, graph valid tree (exactly n - 1 edges and no cycle).
• Redundant connection: return the edge that closes the first cycle.
• Min cost to connect all points and other MST questions via Kruskal.
• Number of islands II: add land cells one by one, union with land neighbors, report the component count each time.
• Accounts merge: union all emails of the same account, then group by root.
• Satisfiability of equality equations, lexicographically smallest equivalent string (keep the smallest letter as the root).
• Smallest string with swaps: union swappable indices, sort characters within each component.
• Evaluate division: weighted DSU with multiplicative offsets.
• Offline edge-limited path queries: sort edges and queries by limit, sweep once.
• Reverse-time processing: when the input removes things (bricks, edges), start from the final state and add them back in reverse order.