Union-Find (Disjoint Set): How It Works, With Python Code

Union-Find, also called a disjoint set union (DSU), keeps track of which elements belong to the same group and lets you merge two groups. It supports two operations, find(x) (which group is x in?) and union(a, b) (merge the groups of a and b), and with two small optimizations each operation takes close to constant time. It is the standard tool for questions like “are these two nodes connected?”, “how many components are there?” and “does adding this edge create a cycle?”.

How it works

Each group is a tree stored in an array. parent[i] is the parent of element i, and the root of a tree (an element that is its own parent) is the group’s representative. Two elements are in the same group if and only if find returns the same root for both.

• find(x) walks up parent links until it reaches the root.
• union(a, b) finds both roots and, if they differ, makes one root the child of the other.

Without any optimization, unioning elements in a bad order creates a long chain, and each find becomes O(n). Two ideas fix that.

Path compression

While walking up to the root in find, point each visited node closer to the root. The next find on those nodes is shorter. The version below uses “path halving”, which does this in a simple loop and avoids the recursion-depth problem of the textbook recursive version.

Union by size (or rank)

When merging, attach the smaller tree under the root of the larger one. This keeps trees shallow. Rank (an upper bound on tree height) works the same way; size is easier to understand and lets you ask for the size of a group for free.

With both optimizations, a sequence of m operations on n elements takes O(m·α(n)) time, where α is the inverse Ackermann function. It grows so slowly that it is at most 4 or 5 for any input that fits in a computer, so treat it as constant.

Implementation

class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))
        self.size = [1] * n
        self.count = n          # number of separate groups

    def find(self, x):
        while self.parent[x] != x:
            self.parent[x] = self.parent[self.parent[x]]   # path halving
            x = self.parent[x]
        return x

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False        # already in the same group
        if self.size[ra] < self.size[rb]:
            ra, rb = rb, ra
        self.parent[rb] = ra    # attach smaller tree under larger
        self.size[ra] += self.size[rb]
        self.count -= 1
        return True

    def connected(self, a, b):
        return self.find(a) == self.find(b)

uf = UnionFind(6)
uf.union(0, 1); uf.union(2, 3); uf.union(1, 2)
print(uf.connected(0, 3))   # True
print(uf.connected(0, 4))   # False
print(uf.count)             # 3  -> groups {0,1,2,3}, {4}, {5}
print(uf.union(0, 3))       # False (already connected)

Returning True or False from union is a small design choice that pays off: many problems need to know whether the union actually merged two different groups.

Union-Find only grows groups. It cannot split a group or remove an element efficiently. If a problem needs deletions, think about processing the operations in reverse (turning deletions into unions) or a different structure.

Use case 1: count connected components

Union every edge, then read count. The same pattern solves “number of provinces” and “number of connected components in an undirected graph”.

def count_components(n, edges):
    uf = UnionFind(n)
    for a, b in edges:
        uf.union(a, b)
    return uf.count

print(count_components(5, [(0, 1), (1, 2), (3, 4)]))   # 2

A DFS or BFS gives the same answer in O(n + e). Union-Find is better when edges arrive one at a time and you need the answer after each one (an “online” problem), because a traversal would have to start over.

Use case 2: detect a cycle in an undirected graph

If an edge connects two nodes that are already in the same group, that edge closes a cycle. This is exactly “Redundant Connection”: return the edge that can be removed to leave a tree.

def find_redundant_connection(edges):
    uf = UnionFind(len(edges) + 1)      # nodes are numbered from 1
    for a, b in edges:
        if not uf.union(a, b):
            return [a, b]

print(find_redundant_connection([(1, 2), (1, 3), (2, 3)]))                  # [2, 3]
print(find_redundant_connection([(1, 2), (2, 3), (3, 4), (1, 4), (1, 5)]))  # [1, 4]

This only works for undirected graphs. For directed graphs, cycle detection needs a DFS with colors or a topological sort.

Use case 3: Kruskal’s minimum spanning tree

Kruskal’s algorithm sorts the edges by weight and adds each edge unless it would form a cycle. Union-Find answers the “would it form a cycle” question in near-constant time, so the sort dominates: O(E log E).

def kruskal(n, edges):
    """edges are (weight, a, b)"""
    uf = UnionFind(n)
    mst, total = [], 0
    for w, a, b in sorted(edges):
        if uf.union(a, b):
            mst.append((a, b, w))
            total += w
    return total, mst

edges = [(1, 0, 1), (2, 1, 2), (3, 2, 3), (4, 3, 0), (5, 0, 2)]
print(kruskal(4, edges))   # (6, [(0, 1, 1), (1, 2, 2), (2, 3, 3)])

Use case 4: grids (number of islands)

A grid is a graph where neighboring land cells are connected. Map cell (r, c) to the index r * C + c, union each land cell with its right and lower land neighbors, then count the groups that contain land.

def num_islands(grid):
    R, C = len(grid), len(grid[0])
    uf = UnionFind(R * C)
    land = 0
    for r in range(R):
        for c in range(C):
            if grid[r][c] == "1":
                land += 1
                for dr, dc in ((1, 0), (0, 1)):
                    nr, nc = r + dr, c + dc
                    if nr < R and nc < C and grid[nr][nc] == "1":
                        uf.union(r * C + c, nr * C + nc)
    return uf.count - (R * C - land)   # drop the one-element groups made by water

print(num_islands(["11000", "11000", "00100", "00011"]))   # 3

For this one-shot problem a DFS is just as good and many people find it easier. Union-Find wins on the follow-up “Number of Islands II”, where land cells are added one at a time and you must report the count after each addition.

Other problems where it fits

• Accounts merge / similar-items grouping: union items that share a key (an email, a label), then group by root.
• Equations like a == b and a != b: union all the equalities first, then check the inequalities.
• Percolation, network connectivity, friend circles: any “are these connected?” with incremental connections.

Complexity

Operation	Time (amortized)
Build with n elements	O(n)
find	O(α(n)), effectively constant
union	O(α(n)), effectively constant

Memory is O(n) for the parent and size arrays.

Common mistakes

• Calling union on the elements instead of on their roots. Set parent[a] = b and you break the structure for every other element in a’s group. Always find the roots first.
• Reading parent[x] directly to compare groups. It is only the root after you call find(x).
• Off-by-one with 1-indexed nodes. Allocate n + 1 elements, as in the redundant connection example.
• Using the recursive find on a long chain in Python. It can hit the recursion limit before path compression has helped. The iterative loop above does not.
• Using it for deletions or directed graphs. Neither is what the structure does.

Practice

Implement the class from memory, then solve Number of Provinces, Redundant Connection and Number of Islands II. If graph traversal is still shaky, read breadth-first search first, since Union-Find is usually the second tool you reach for on connectivity problems. You can practice the graph side on AlgoCademy with the Find if Path Exists in Graph exercise, and see when to use a visited set for the traversal alternative.