Union-Find (DSU)

Union-Find (DSU) | Algoryva

First, the naive way

Label everyone with their group number

The direct approach keeps an array label[i] saying which group item i is in. Checking whether two items are connected is instant: compare labels.

But merging two groups means relabelling every member of one of them. After many merges, nearly everyone may be in one huge group, so each merge touches about n items.

python

class Labels:
    def __init__(self, n):
        self.label = list(range(n))
    def connected(self, a, b):
        return self.label[a] == self.label[b]
    def union(self, a, b):
        old, new = self.label[a], self.label[b]
        for i in range(len(self.label)):      # touches every item
            if self.label[i] == old:
                self.label[i] = new

L = Labels(5)
L.union(0, 1)
L.union(1, 2)
print(L.connected(0, 2))   # True
print(L.connected(0, 4))   # False

Cost: O(1) to ask, but O(n) per merge, so O(n × q) for q merges. Too slow for n, q around 10^5 or more.

Worked by hand

The same six merges by hand

Start with parent = [0, 1, 2, 3, 4, 5] (everyone alone, all sizes 1). Ties in size go to the first root.

operation	roots found	action	parent afterwards
union(0, 1)	0, 1	attach 1 under 0	[0, 0, 2, 3, 4, 5]
union(2, 3)	2, 3	attach 3 under 2	[0, 0, 2, 2, 4, 5]
union(1, 3)	0, 2	sizes equal (2, 2): attach 2 under 0	[0, 0, 0, 2, 4, 5]
union(4, 5)	4, 5	attach 5 under 4	[0, 0, 0, 2, 4, 4]
union(0, 2)	0, 0	same root: already together, do nothing	[0, 0, 0, 2, 4, 4]
union(3, 5)	0 (via 3, 2), 4	sizes 4 and 2: attach 4 under 0; path compression sets parent[3] = 0	[0, 0, 0, 0, 0, 4]

1
The no-op step matters
union(0, 2) found equal roots. In graph terms that edge would close a cycle. This is exactly how Kruskal's algorithm and “redundant edge” problems use Union-Find.
2
Path compression at work
To find 3's root the code walked 3, 2, 0, then rewired parent[3] to 0. The next find(3) takes one hop.

Implementation

Union-Find in Python

Both optimisations fit in a few lines. This class is worth memorising; it appears in many contest solutions.

Plan in plain steps

pseudocode

find(x): walk to the root; then re-walk, pointing every node at the root
union(a, b): ra, rb = find(a), find(b)
    if same, stop (return False)
    attach the smaller tree under the larger; add sizes; return True

python

class DSU:
    def __init__(self, n):
        self.parent = list(range(n))
        self.size = [1] * n
        self.groups = n

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

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False
        if self.size[ra] < self.size[rb]:
            ra, rb = rb, ra
        self.parent[rb] = ra                   # smaller under larger
        self.size[ra] += self.size[rb]
        self.groups -= 1
        return True

d = DSU(6)
for a, b in [(0, 1), (2, 3), (1, 3), (4, 5)]:
    d.union(a, b)
print(d.find(3) == d.find(0))   # True
print(d.find(3) == d.find(4))   # False
print(d.groups)   # 2
print(d.union(0, 2))   # False

Reading the code

›find makes two passes: the first locates the root; the second rewires each visited item to point at it. After that, every item on the path is one hop from the root.
›union returns True if it merged two different groups and False if they were already together. That return value is what makes cycle detection a one-liner.
›The if self.size[ra] < self.size[rb] swap ensures the larger tree stays as the root.
›groups counts remaining components: it starts at n and drops by one per successful merge.

Worked problem

Count the provinces

The problem

Given an n × n matrix where m[i][j] = 1 means city i and city j are directly connected, count the groups of cities that are connected directly or indirectly.

Start with n groups, one per city.
For every pair (i, j) that is directly connected, union them. Each successful union merges two groups, so subtract one.
What is left is the number of provinces.

python

def find(p, x):
    while p[x] != x:
        p[x] = p[p[x]]        # path halving: point at the grandparent
        x = p[x]
    return x

def provinces(m):
    n = len(m)
    p = list(range(n))
    groups = n
    for i in range(n):
        for j in range(i + 1, n):
            if m[i][j]:
                ri, rj = find(p, i), find(p, j)
                if ri != rj:
                    p[ri] = rj
                    groups -= 1
    return groups

print(provinces([[1, 1, 0], [1, 1, 0], [0, 0, 1]]))   # 2

Takeaway. “Count the groups” can be solved by DFS or by Union-Find. Union-Find wins when the connections arrive incrementally. This snippet also shows *path halving*, a compact variant of compression that needs no second pass.

Worked problem

Find the redundant connection

The problem

A tree with nodes 1 to n had one extra edge added, making a cycle. Given the edge list in order, return the extra edge that closes the cycle (the last such edge in the list).

Process edges in order.
For each edge, find the roots of its two endpoints.
If the roots are equal, the endpoints were already connected, so this edge closes a cycle: return it.
Otherwise union them.

python

def find(p, x):
    while p[x] != x:
        p[x] = p[p[x]]
        x = p[x]
    return x

def redundant(edges):
    p = list(range(len(edges) + 1))      # nodes are 1..n, so size n + 1
    for a, b in edges:
        ra, rb = find(p, a), find(p, b)
        if ra == rb:
            return [a, b]
        p[ra] = rb
    return []

print(redundant([[1, 2], [1, 3], [2, 3]]))   # [2, 3]

Takeaway. An edge whose endpoints already share a root closes a cycle. This is the same test Kruskal's algorithm uses to skip edges that would create a cycle when building a minimum spanning tree.

Wrap-up

Remember this

Union-Find tracks groups that only ever merge, answering “same group?” in near-constant time.
Each group is a tree in one parent array; the root identifies the group.
Union by size keeps trees shallow; path compression flattens them as you search.
A union between items that already share a root means a cycle, the key to Kruskal's algorithm.
Use it when connections arrive over time; use DFS/BFS when the graph is fixed.

Cheat sheet

find(x): follow parent to the root; compress the path
union(a, b): link the smaller root under the larger; False if same root
Connected?: find(a) == find(b)
Cost: ≈ O(1) each: O(α(n)) amortised
Components: start at n, subtract 1 per successful union
Cannot: split groups, or return paths

Top traps to avoid

Comparing parent[a] == parent[b] instead of the roots.
Linking items instead of roots (parent[a] = b).
Skipping the size (or rank) rule.

Mastery check

0 of 15 answered

Answer from understanding, not memory of the page. You need 12 of 15. Each answer is explained, and you can retry as often as you like.

1.In Union-Find, how do you test whether items a and b are in the same group?
2.After union(0, 1) and union(1, 2) with DSU(3) (equal sizes go to the first root), how many groups remain?
3.What does path compression do?
4.With union by size and path compression, what is the cost per operation in practice?
5.Which problem is the best fit for Union-Find?
6.union(a, b) returns False. What does that tell you?
7.p = list(range(5)); union(0,1) and union(1,2) with p[find(a)] = find(b). What does find(0) return?
8.Four nodes with unions (0,1) and (2,3). How many separate sets are there?
9.Three nodes, after p[find(0)] = find(1). What is printed by print(find(0) == find(1), find(0) == find(2))?
10.Kruskal's cycle check: unions (0,1) then (1,0). How many edges are redundant (joining nodes already in one set)?
11.Which two optimisations make union-find almost constant time per operation?
12.With path compression and union by rank, the time per operation is roughly:
13.Six nodes, unions (0,1), (2,3), (4,5), each joining different sets. How many sets remain?
14.Which question does union-find answer best?
15.Path p = [1, 2, 3, 4, 4]. After find(0) with full path compression, what is p[0]?

Union-Find (DSU)

After this lesson you will be able to

Builds on

Unlocks

Who is in the same friend group?

Label everyone with their group number

Store a pointer to the leader, not a copy of it

A forest in one array

Merging groups step by step

The same six merges by hand

Union-Find in Python

Reading the code

Why it is nearly constant

Edge cases

Mistakes to avoid

How to spot a Union-Find problem

Count the provinces

Find the redundant connection

Wrap-up

Remember this

Cheat sheet

Top traps to avoid

Mastery check

Practice this topic

Warm-up2 problems

Core8 problems

Application30 problems

Challenge20 problems