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

Recommended DSA resources

Some links are affiliate links. If you buy through them, Laviyx Learn may earn a commission at no extra cost to you. As an Amazon Associate, we earn from qualifying purchases.