Skip to content
Laviyx Learn

Try

DSA course menu

Union-Find (Disjoint Set Union) in Python

Learn the union-find data structure: find and union operations, path compression and union by size, nearly O(1) operations, counting connected components and detecting cycles in undirected graphs.

Advanced10 min readUpdated
On this page

Union-find (also called disjoint set union, DSU) keeps track of items split into groups, and answers two questions very fast:

  • find(x): which group is x in?
  • union(x, y): merge the groups of x and y.

Two items are in the same group exactly when find returns the same answer for both. It is ideal when connections only ever get added: friends of friends, merging accounts, pixels joining into regions, and Kruskal's algorithm.

The idea: a forest of parent pointers#

Each group is a tree. Every item points to a parent; the root of the tree points to itself and represents the whole group. find follows parents up to the root; union hangs one root under the other.

Python
class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))     # everyone starts alone

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

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra != rb:
            self.parent[rb] = ra

uf = UnionFind(6)
uf.union(0, 1)
uf.union(1, 2)
uf.union(3, 4)
print(uf.parent)
print(uf.find(2) == uf.find(0), uf.find(3) == uf.find(0), uf.find(5))
Output
[0, 0, 0, 3, 3, 5]
True False 5

This works, but trees can grow tall, and then find walks a long chain.

Two optimisations#

Union by size: always hang the smaller tree under the root of the larger one, so trees stay shallow.

Path compression: after find reaches the root, point every node on the way directly at the root, so the next find is one step.

Python
class UnionFind:
    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
            self.parent[x], x = root, self.parent[x]
        return root

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

uf = UnionFind(8)
for a, b in [(0, 1), (2, 3), (1, 3), (4, 5), (6, 7), (5, 7)]:
    uf.union(a, b)
print("groups:", uf.groups)
print("size of 0's group:", uf.size[uf.find(0)])
print(uf.find(3) == uf.find(0), uf.find(4) == uf.find(0))
Output
groups: 2
size of 0's group: 4
True False

With both optimisations, each operation takes O(α(n)) amortized time, where α is the inverse Ackermann function. It grows so slowly that it is below 5 for any input that could ever exist, so in practice every operation is constant time.

Detecting a cycle in an undirected graph#

Add edges one at a time. If an edge joins two vertices that are already in the same group, there was already a path between them, so the edge closes a cycle:

Python
class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))
    def find(self, x):
        while self.parent[x] != x:
            self.parent[x] = self.parent[self.parent[x]]   # halve the path
            x = self.parent[x]
        return x
    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False
        self.parent[rb] = ra
        return True

def first_cycle_edge(n, edges):
    uf = UnionFind(n)
    for a, b in edges:
        if not uf.union(a, b):
            return (a, b)
    return None

print(first_cycle_edge(4, [(0, 1), (1, 2), (2, 3)]))
print(first_cycle_edge(4, [(0, 1), (1, 2), (2, 0), (2, 3)]))
Output
None
(2, 0)

This version uses path halving (each node skips to its grandparent), a simpler form of path compression with the same guarantees.

Union-find or DFS?#

Both find connected components. Union-find wins when edges arrive over time and you must answer "are these connected?" between arrivals; DFS would have to search again each time. Union-find cannot handle edges being removed.

Common mistakes#

Practice#

Practice exercise

Accounts merge (simplified): each person entry is a list of email addresses. Two entries belong to the same person if they share any email. Write count_people(entries) that returns how many distinct people there are.

Show hint

Give each entry an index. Map each email to the first entry that had it; when an email is seen again, union the two entries. Count the groups at the end.

Show solution
Python
class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))
        self.groups = n
    def find(self, x):
        while self.parent[x] != x:
            self.parent[x] = self.parent[self.parent[x]]
            x = self.parent[x]
        return x
    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra != rb:
            self.parent[rb] = ra
            self.groups -= 1

def count_people(entries):
    uf = UnionFind(len(entries))
    owner = {}
    for i, emails in enumerate(entries):
        for e in emails:
            if e in owner:
                uf.union(owner[e], i)
            else:
                owner[e] = i
    return uf.groups

entries = [
    ["asha@a.com", "asha@work.com"],
    ["ravi@r.com"],
    ["asha@work.com", "asha.k@b.com"],
    ["meera@m.com"],
    ["asha.k@b.com"],
]
print(count_people(entries))
Output
3

Entries 0, 2 and 4 are linked through shared emails, so there are 3 people.

Key takeaway#

Next: connecting everything as cheaply as possible, with Minimum Spanning Trees.

Finished this lesson?

Related lessons

DSAMinimum Spanning Trees: Kruskal's and Prim's Algorithms in PythonLearn minimum spanning trees: connecting every vertex of a weighted graph at the lowest total cost, Kruskal's algorithm with union-find, Prim's algorithm with a heap, and real uses like network design.Advanced8 min readDSADepth-First Search (DFS) in Python: Recursive and IterativeLearn depth-first search: recursive and stack-based DFS, connected components, counting islands in a grid, detecting cycles in undirected and directed graphs, and when to choose DFS over BFS.Intermediate11 min readDSAGraphs in Data Structures: Terminology and Representation in PythonLearn graph basics: vertices and edges, directed and undirected, weighted graphs, degree, paths and cycles, and how to store a graph as an adjacency list, adjacency matrix or edge list in Python.Intermediate9 min read