Disjoint Set Union (union-find)

Some problems ask you to track which items belong together. Friends in a social network form friend groups. Cities connected by roads form regions. Rock samples from the same deposit form clusters. You start with each item alone, then learn that pairs belong together, and need to answer whether two items are in the same group.

A disjoint set uniondisjoint set unionA data structure that tracks which elements belong to the same group, supporting fast union and connectivity queries.In the glossary data structure, also called union-findunion-findAnother name for disjoint set union; commonly used when describing the find and union operations.In the glossary, tracks these groups efficiently. You can ask "are A and B in the same group?" and merge two groups when you learn they are connected. Both operations run in nearly constant time.

Find and union with trees

The simplest idea is to use a parent array. Each element points to its parent. An element with no parent is a root. All elements in the same tree have the same root, so they are in the same group.

To check if A and B are in the same group, follow parents up from A to its root and from B to its root. If the roots match, they are in the same group. The operation is called find.

To merge the group containing A with the group containing B, find their roots. Then make one root a child of the other. The operation is called union.

Here is a simple example. Start with five independent elements.

1def find(parent, x):2    """Follow parent pointers to the root."""3    while parent[x] != x:4        x = parent[x]5    return x6
7
8def union(parent, a, b):9    """Merge the groups containing a and b."""10    root_a = find(parent, a)11    root_b = find(parent, b)12    if root_a != root_b:13        parent[root_a] = root_b14
15
16# Start with 5 independent elements17parent = [0, 1, 2, 3, 4]18
19# Union 0 and 120union(parent, 0, 1)21print("After union(0, 1):", parent)22
23# Union 2 and 324union(parent, 2, 3)25print("After union(2, 3):", parent)26
27# Check connectivity28print("find(0):", find(parent, 0))29print("find(1):", find(parent, 1))30print("find(2):", find(parent, 2))31print("find(3):", find(parent, 3))

After union(0, 1): [1, 1, 2, 3, 4]
After union(2, 3): [1, 1, 3, 3, 4]
find(0): 1
find(1): 1
find(2): 3
find(3): 3

This code creates a parent array where each element is its own parent initially. Both find and union work.

Path halving keeps trees shallow

A problem with trees is that they can become very deep. If you merge by always making the same root a child of another, you build a chain. Then find on the deepest element takes linear time. Doing this a million times is far too slow.

The fix is to flatten trees as you traverse them. When you follow a parent pointer from node u to its parent parent[u], also point u to its grandparent parent[parent[u]]. This is called path halvingpath halvingAn optimization in the find operation where each node points to its grandparent instead of its parent, flattening the tree over time.In the glossary. Over many finds, the trees stay shallow.

1def find(parent, x):2    """Find root with path halving."""3    while parent[x] != x:4        parent[x] = parent[parent[x]]  # Skip one level5        x = parent[x]6    return x7
8
9def union(parent, a, b):10    """Merge the groups containing a and b."""11    root_a = find(parent, a)12    root_b = find(parent, b)13    if root_a != root_b:14        parent[root_a] = root_b15
16
17# Build a chain by naive unions: 0 <- 1 <- 2 <- 3 <- 418parent = list(range(5))19for i in range(4):20    union(parent, i, i + 1)21
22print("Parent array after naive unions:", parent)23
24# First find on 0: takes 4 steps without halving25print("find(0) with path halving:", find(parent, 0))26
27# After halving, the path is shorter28print("Parent array after find:", parent)

Parent array after naive unions: [1, 2, 3, 4, 4]
find(0) with path halving: 4
Parent array after find: [2, 2, 4, 4, 4]

The find function now climbs to the root in fewer steps. Each time you call find on an element, the path gets shorter.

Union by size prevents imbalance

Another way to keep trees shallow is to always attach the smaller tree under the larger tree. When you merge two groups, find their roots and attach the root of the smaller tree under the root of the larger tree. This is called union by sizeunion by sizeAn optimization where the smaller tree is always attached under the larger tree during union, keeping the forest shallow.In the glossary.

Combining path halving with union by size makes each operation run in nearly constant time. After any sequence of operations, a find or union takes time proportional to the logarithm of the number of elements, or often much faster in practice.

1def find(parent, x):2    """Find root with path halving."""3    while parent[x] != x:4        parent[x] = parent[parent[x]]5        x = parent[x]6    return x7
8
9def union(parent, size, a, b):10    """Merge groups containing a and b; attach smaller tree under larger."""11    root_a = find(parent, a)12    root_b = find(parent, b)13    if root_a == root_b:14        return15    # Attach smaller tree under larger16    if size[root_a] < size[root_b]:17        parent[root_a] = root_b18        size[root_b] += size[root_a]19    else:20        parent[root_b] = root_a21        size[root_a] += size[root_b]22
23
24# Start with 5 elements25n = 526parent = list(range(n))27size = [1] * n28
29# Union 0 and 130union(parent, size, 0, 1)31print("After union(0, 1):", parent)32
33# Union 2 and 334union(parent, size, 2, 3)35print("After union(2, 3):", parent)36
37# Union the two groups38union(parent, size, 0, 2)39print("After union(0, 2):", parent)40print("Sizes:", size)

After union(0, 1): [0, 0, 2, 3, 4]
After union(2, 3): [0, 0, 2, 2, 4]
After union(0, 2): [0, 0, 0, 2, 4]
Sizes: [4, 1, 2, 1, 1]

The size array tracks how many elements are in each tree. When you merge two roots, you attach the smaller one under the larger and update the size.

Connectivity in a network

Imagine a system where computers are initially isolated. You learn about five new network cables, each connecting two computers. Your job is to report when two computers are connected, even indirectly through other machines.

This is a natural use case for union-find. You start with each computer as its own group. Then you process each cable by calling union on its two endpoints. Afterward, you can query whether two computers are connected by calling find and comparing roots.

1import sys2
3
4def main() -> None:5    input_data = sys.stdin.read().split()6    if not input_data:7        return8
9    pos = 010    n = int(input_data[pos]); pos += 111    m = int(input_data[pos]); pos += 112
13    parent = list(range(n))14    size = [1] * n15
16    def find(x: int) -> int:17        while parent[x] != x:18            parent[x] = parent[parent[x]]19            x = parent[x]20        return x21
22    def union(a: int, b: int) -> None:23        root_a = find(a)24        root_b = find(b)25        if root_a == root_b:26            return27        if size[root_a] < size[root_b]:28            parent[root_a] = root_b29            size[root_b] += size[root_a]30        else:31            parent[root_b] = root_a32            size[root_a] += size[root_b]33
34    for _ in range(m):35        u = int(input_data[pos]); pos += 136        v = int(input_data[pos]); pos += 137        union(u, v)38
39    q = int(input_data[pos]); pos += 140    out_lines = []41    for _ in range(q):42        a = int(input_data[pos]); pos += 143        b = int(input_data[pos]); pos += 144        out_lines.append("yes" if find(a) == find(b) else "no")45
46    print("\n".join(out_lines))47
48
49if __name__ == "__main__":50    main()

5 3
0 1
1 2
3 4
3
0 2
0 3
1 3

yes
no
no

The input gives the number of computers and cables, then one line per cable with two computer ids, then the number of queries and one line per query with two computer ids. The program prints yes or no for each query. Union-find answers each query instantly once all the cables have been processed.

The next free slot trick

A clever optimization applies DSU in a different way. Suppose you have a list of parking slots and many parked cars. When a car leaves, its slot becomes free. When a new car arrives, it wants to park in the lowest-numbered free slot. A brute-force search takes linear time per arrival. With DSU, you can answer in nearly constant time.

The trick is to use DSU on the slots themselves. When a car parks in slot s, you union s with s + 1. This creates a tree rooted at the first free slot after s. When the next car asks for the lowest free slot, you start at slot 1 and call find(1). If slot 1 is full, find(1) jumps past it to the root of its tree, which points to the first free slot.

1def find(parent, x):2    """Find root with path halving; returns first free slot at or after x."""3    if parent[x] == x:4        return x5    parent[x] = find(parent, x + 1)6    return parent[x]7
8
9def main() -> None:10    n = 511    parent = list(range(n + 1))12
13    results = []14    for _ in range(3):15        # Find next free slot16        slot = find(parent, 1)17        results.append(str(slot))18        # Mark it occupied by pointing it to the next slot19        if slot < n:20            parent[slot] = slot + 121
22    print("\n".join(results))23
24
25if __name__ == "__main__":26    main()

1
2
3

Here you have parking slots 1 through N. Each car arrives and parks in the lowest free slot, then leaves after one turn. The output is which slot each car used. Union-find skips over full slots in one operation instead of searching sequentially.

Common mistakes

A common mistake is using recursion for find. If the tree becomes deep before path halving takes effect, a recursive find hits Python's recursion limit. Always use a loop.

Another error is forgetting to update the size when you union. If you attach one root under another without updating the size of the parent, future unions may attach the larger tree under the smaller one, defeating the purpose.

Avoid building an adjacency list and using BFS just to check connectivity. Union-find handles that exact question with less code and less time.