Union-Find / Disjoint Set Basics: Parent Array, Find, Union

Union-find, also called a disjoint set structure, answers one question very cheaply: are these two things in the same group? Groups only ever merge, and the structure needs nothing more than one array and two short functions. This article is about those basics: what the array means, how find and union work on it, and why it beats re-running a graph search when merges and questions arrive mixed together.

The problem it solves

Many problems are "groups that merge over time":

• People become friends; are two people in the same friend circle?
• Two words are declared synonyms; are two other words equivalent, through any chain of synonyms?
• Two computers are cabled together; can they reach each other?
• Two photo clusters are confirmed to be the same person; which photos belong together?

You could keep an explicit list per group and copy members on every merge, but merges then cost the size of a group. You could build a graph and run BFS for every question, but each question then costs the whole graph. Union-find makes both operations cheap, and it handles the online case, where merges and questions are interleaved, without rebuilding anything.

The intuition

Keep one array, parent. parent[x] is the element x points at. At the start, every element points at itself: n separate groups.

• The root is the group's name. Follow pointers from any element and you stop at an element that points at itself, parent[r] == r. That root names the whole group.
• find(x) walks up from x to its root.
• same(a, b) is find(a) == find(b).
• union(a, b) finds both roots. Equal roots mean they are already one group. Otherwise, point one root at the other: a single write merges two whole groups, however large.

Nothing ever lists a group's members, and nothing moves when groups merge. The trees can be shaped differently depending on the order of unions, but the partition into groups is the same.

The basic version has one weakness: unions can build a long chain, and find then walks it one step at a time. Path compression and union by size fix that, and they are the subject of the path compression article.

Watch it run

The animation runs the lesson's own code on six elements, each its own parent: six separate groups before anything is joined, parent array 0 to 5. union(0, 1): both are their own root, so one root is pointed at the other, and parent[0] becomes 1. union(1, 2): 1 is still a root, so it hangs under 2. One write, three elements joined, and the array reads 1, 2, 2, 3, 4, 5. Then find(0) starts at 0: parent[0] is not 0, so it follows the pointer. At 1, parent[1] is 2, not 1, so it keeps walking up. parent[2] is 2, so the walk stops: 2 is the root, and the root is the group's name. find(2) is 2 as well; equal roots, so 0 and 2 are in one set. find(3) is 3, a different root, so a different set. No member list was ever read.

The code

The lesson's code, printing the parent array after each union, the two questions from the animation, and the groups read back by root:

import random

parent = list(range(6))          # everyone starts alone: parent[i] == i

def find(x):
    while parent[x] != x:        # walk up to the root
        x = parent[x]
    return x

def union(a, b):
    ra, rb = find(a), find(b)
    if ra == rb:
        return False             # already one group
    parent[ra] = rb              # hang one root under the other
    return True

union(0, 1)
print(parent)                    # [1, 1, 2, 3, 4, 5]
union(1, 2)
print(parent)                    # [1, 2, 2, 3, 4, 5]
print(find(0) == find(2), find(0) == find(3), union(2, 0))   # True False False

def groups(n):
    out = {}
    for x in range(n):
        out.setdefault(find(x), []).append(x)
    return sorted(out.values())

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

union(2, 0) returns False: they already share a root. Real problems use names, not indices, so the same idea with a dictionary, where any new label starts as its own root. Synonym pairs arrive, questions are asked between them, and a later pair joins two groups without anything being rebuilt. Then the weakness: a thousand unions in a line build a chain, and find(0) walks all of it:

class DisjointSets:
    """The same two operations over any hashable labels, added on first sight."""
    def __init__(self):
        self.parent = {}
    def find(self, x):
        self.parent.setdefault(x, 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[ra] = rb
        return ra != rb
    def same(self, a, b):
        return self.find(a) == self.find(b)

words = DisjointSets()
for a, b in [("fast", "quick"), ("quick", "rapid"), ("big", "large")]:
    words.union(a, b)
print(words.same("fast", "rapid"), words.same("fast", "large"))   # True False
words.union("rapid", "large")                                      # a new pair arrives later
print(words.same("fast", "big"))                                   # True

def hops_to_root(x):
    steps = 0
    while parent[x] != x:
        x, steps = parent[x], steps + 1
    return steps

parent = list(range(1000))
for i in range(999):
    union(i, i + 1)                    # each union hangs the old root one level lower
print(hops_to_root(0), hops_to_root(999))   # 999 0

Every union in that loop was cheap, but the tree became a path of 1,000 elements, so one find costs 999 steps. That is the case the optimisations exist for.

The seeded check runs 500 random sequences of unions and questions against a brute force that keeps explicit Python sets and merges them by copying. Every union result and every same answer must match, and at the end the groups read back by root must be exactly the brute-force sets:

ok = True
for seed in range(500):
    r = random.Random(seed)
    n = r.randint(1, 30)
    ds, sets = DisjointSets(), [{x} for x in range(n)]          # brute force: explicit sets
    for _ in range(r.randint(0, 60)):
        a, b = r.randrange(n), r.randrange(n)
        sa = next(s for s in sets if a in s)
        sb = next(s for s in sets if b in s)
        if r.random() < 0.5:
            ok &= ds.union(a, b) == (sa is not sb)
            if sa is not sb:
                sets.remove(sa)
                sb |= sa
        else:
            ok &= ds.same(a, b) == (sa is sb)
    by_root = {}
    for x in range(n):
        by_root.setdefault(ds.find(x), set()).add(x)
    ok &= sorted(map(sorted, by_root.values())) == sorted(map(sorted, sets))
    ok &= all(ds.parent.get(root, root) == root for root in by_root)
print(ok)                                                             # True

The complexity

• Space: O(n), one parent entry per element.
• find: O(depth), which is O(n) in the worst case for this basic version, as the chain shows.
• union: two finds plus one write.
• With path compression and union by size: nearly constant amortised per operation; see the Big-O cheat sheet.

Where it goes wrong

• Pointing elements instead of roots. parent[a] = b without finding roots first can detach part of a group. Always link root to root.
• Comparing parents instead of roots. parent[a] == parent[b] is not "same group"; compare find results.
• Expecting deletes. Union-find only merges; splitting a group needs a rebuild.
• Recursion in find. A recursive version hits the recursion limit on a long chain; keep it a loop.

When it shows up in interviews

As friend circles, accounts merge, synonym or equation equivalence, "number of provinces", and as the cycle check inside Kruskal's algorithm. Two follow-ups are common: counting groups, covered in connected components, and detecting the edge that closes a cycle, in redundant connection.

How to say it in an interview

"Union-find keeps a parent array where every element points towards its group's root, and a root points at itself, so the root is the group's name. Find walks up to the root; two elements are in the same group exactly when their roots match. Union finds both roots and points one at the other, merging two groups with one write. Unoptimised, find can be O(n) on a chain, so in practice I add path compression and union by size, which make operations nearly constant amortised."