Union-find

Problem

The input is a collection of items and a stream of "these two belong together" facts: edges of a graph, accounts that share an email, numbers that share a factor. You are asked how many groups there are, whether a new fact was redundant, whether everything ended up in one group, or what the groups are. The answer must not require rebuilding anything when the next fact arrives.

Intuition

Picture friend groups at a party. Each group has a spokesperson. To check whether two people are in the same group, ask each of them who their spokesperson is, passing the question up the line of "my friend X introduced me" until someone says "that's me". If the answers match, they are together. To join two groups, make one spokesperson report to the other. That is the whole data structure: parent[x] is who introduced x, and the root is its own parent.

Left alone, the chains of introductions can grow long, and asking becomes slow. Two cheap habits fix it. When you join two groups, hang the smaller one under the larger one's spokesperson, so no chain gets deeper unless the group at least doubles. And whenever you walk up a chain, repoint everyone you passed straight at the spokesperson, so the next question is one step. Together they make an operation cost the inverse Ackermann function of n, which is below five for any input that fits in a universe.

The same unions, run in the same order, can build a chain or a flat star depending on which root goes under which. The left tree makes a find from node 0 cost 7 steps; the right one costs 1. The diagram scrolls sideways on a narrow screen.

Mechanics

Find and union

find(x) follows parent pointers until a node is its own parent and returns it. union(a, b) finds both roots; if they are equal the pair was already connected (report it, that is the redundant edge), and otherwise one root is pointed at the other and the group count drops by one.

Invariants

Two facts hold after every operation. First, following parents from any node ends at a root and never loops, because a union only points a root at a different root. Second, two nodes share a root exactly when a chain of union calls connected them: a union merges two whole groups and nothing else changes membership. So find(a) == find(b) is the same as "a and b are connected by the edges seen so far".

Why it is fast

With union by size, a node's depth increases only when its group is attached under a group at least as large, which at least doubles the group it belongs to, so the depth is at most log₂ n. Path compression flattens whatever a find walked. Used together, Tarjan proved the amortized cost per operation is O(α(n)), the inverse Ackermann function. Either one alone already gives O(log n) amortized, which is plenty for interview limits.

Counting and redundancy

Keep components = n and subtract one on each successful union. A union that finds equal roots is a cycle edge in an undirected graph. A tree on n nodes has exactly n - 1 edges and a single component, so it is valid exactly when the edge count is n - 1 and no edge was redundant.

Blue nodes are roots. Nodes 0, 1, 2 and 3 share root 0; nodes 4 and 5 share root 4; node 6 is alone.

Worked example

Seven nodes and the edges 0-1, 2-3, 1-2, 4-5, 3-0 in that order. Marking each as merged or redundant: 0-1 merges; 2-3 merges; 1-2 merges; 4-5 merges; 3-0 is redundant. The final parent array is 0, 0, 0, 0, 4, 4, 6 and the count is 3 groups. The first redundant edge is 3-0, which is the answer to "which edge closes a cycle".

Why the two habits matter. Eight nodes, the seven unions 1-0, 2-1, 3-2, 4-3, 5-4, 6-5, 7-6, then three repeated queries of the pair 0 and 7:

The unlucky union order builds a chain, so without union by size a single find walks 7 steps; compression cuts the repeated queries down but not the first expensive one. Union by size prevents the chain.

Is it a tree? Five nodes with edges 0-1, 0-2, 0-3, 1-4: the count is 4 = 5 - 1 and no edge is redundant, so the answer is True. Replace the last edge so the edges are 0-1, 1-2, 2-3, 1-3: the count is still 4 but edge 1-3 closes the loop 1-2-3 and leaves node 4 isolated, so it is False.

Merging accounts. Treat each email as a node, union all emails inside one account, and group by root. The 4 accounts (two of the Johns share a@m) merge to John a@m b@m d@m; John c@m; Mary e@m.

Implementation

class DSU:
    def __init__(self, n):
        self.parent = list(range(n))
        self.size = [1] * n
        self.components = 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:
            return False
        if self.size[ra] < self.size[rb]:
            ra, rb = rb, ra
        self.parent[rb] = ra
        self.size[ra] += self.size[rb]
        self.components -= 1
        return True


def count_components(n, edges):
    dsu = DSU(n)
    for a, b in edges:
        dsu.union(a, b)
    return dsu.components


def find_redundant_connection(edges):
    dsu = DSU(len(edges) + 1)
    for a, b in edges:
        if not dsu.union(a, b):
            return [a, b]
    return []

Variations and edge cases

def valid_tree(n, edges):
    if len(edges) != n - 1:
        return False
    dsu = DSU(n)
    return all(dsu.union(a, b) for a, b in edges)


def accounts_merge(accounts):
    index = {}
    dsu = DSU(len(accounts))
    for i, (_, *emails) in enumerate(accounts):
        for email in emails:
            if email in index:
                dsu.union(i, index[email])
            else:
                index[email] = i
    groups = {}
    for email, i in index.items():
        groups.setdefault(dsu.find(i), []).append(email)
    return [[accounts[root][0]] + sorted(emails) for root, emails in groups.items()]


def can_traverse_all(nums):
    n = len(nums)
    if n == 1:
        return True
    dsu = {}

    def find(x):
        dsu.setdefault(x, x)
        while dsu[x] != x:
            dsu[x] = dsu[dsu[x]]
            x = dsu[x]
        return x

    for i, v in enumerate(nums):
        if v == 1:
            return False
        p = 2
        while p * p <= v:
            if v % p == 0:
                dsu[find(("p", p))] = find(("i", i))
                while v % p == 0:
                    v //= p
            p += 1
        if v > 1:
            dsu[find(("p", v))] = find(("i", i))
    return len({find(("i", i)) for i in range(n)}) == 1

• Compare roots, not parents. After unions a node's parent may be an interior node. Always call find before comparing two nodes.
• Recursion. A recursive find on a long chain overflows Python's stack. The iterative path-halving version above never does.
• Keys that are not 0..n-1. Use a dictionary for the parent map (emails, strings, tuples) or map the keys to integers first, as the account example does.
• Grids. Number cells r * cols + c. For "add land, report islands" streams, union only with neighbors that are already land, and update the count on each add.
• Directed edges. Plain union-find ignores direction, so it cannot detect a cycle in a directed graph; use a three-state DFS or Kahn's algorithm there. A directed variant of Redundant Connection needs extra care for nodes with two parents.
• Edge cases: a self-loop (the root is always equal, so it is redundant), duplicate edges, zero edges, n of 1 (a lone node is a tree), and nodes that never appear in an edge but still count as components.

Complexity and tradeoffs

• Time: O(α(n)) amortized per find or union with both optimizations, so O(E · α(n)) for E edges, in practice linear. With one optimization it is O(log n) amortized.
• Space: O(n) for the parent array, plus O(n) for sizes. Dictionaries add a constant factor.
• Baseline. Re-running a DFS after every added edge costs O(E · (V + E)). Union-find processes the same stream in near-linear time.
• Versus DFS/BFS. For one static graph, a single traversal counts components just as fast and needs no new structure. Union-find wins when edges arrive over time, when you need to detect the first redundant edge, or when the graph is implicit in a rule (shared factors, shared emails).
• When it fails. It cannot delete an edge, answer path or distance questions, or use edge direction. Offline deletions can sometimes be reversed into additions processed backwards. Minimum spanning trees use union-find inside Kruskal's algorithm.

Problems

Number of Connected Components and Graph Valid Tree are the direct counts: components after all unions, and the n - 1 edge rule. Redundant Connection returns the first edge whose endpoints already share a root. Accounts Merge and Greatest Common Divisor Traversal hide the edges: shared emails, shared prime factors. Several problems homed elsewhere also fall to a union-find on a grid or a sorted edge list.

Medium

• 261. Graph Valid Tree Blind 75 Grind 169 NeetCode 150 NeetCode 250
  A graph is a tree when it has exactly n-1 edges and no edge ever joins two nodes that are already connected.
• 323. Number of Connected Components in an Undirected Graph Blind 75 Grind 169 NeetCode 150 NeetCode 250
  Start with n components and subtract one every time a union joins two different sets.
• 684. Redundant Connection NeetCode 150 NeetCode 250
  Union the endpoints of each edge in order; the first edge whose endpoints are already connected closes the cycle.
• 721. Accounts Merge Grind 169 NeetCode 250
  Union every email in an account with the account's first email, then group emails by root and sort each group.

Hard

• 2709. Greatest Common Divisor Traversal NeetCode 250
  Union each index with the prime factors of its value, then check that all indices end up in one set.

Also uses this pattern

These problems have a different home page, but this idea is part of their solution.

• 200. Number of Islands Medium
• 399. Evaluate Division Medium
• 1631. Path With Minimum Effort Medium
• 778. Swim in Rising Water Hard
• 1489. Find Critical and Pseudo-Critical Edges in Minimum Spanning Tree Hard

Connections

Union-find is graph traversal's answer when the graph grows. It is also the cycle test inside Kruskal's minimum spanning tree algorithm, and with edges sorted by cost it solves "smallest bottleneck" questions that look like shortest-path problems. Grid problems with changing land use the same structure with numbered cells.

Check yourself

References

• Cormen, Leiserson, Rivest and Stein, Introduction to Algorithms, chapter 19 (disjoint-set forests, union by rank and path compression, with the analysis).
• Tarjan, "Efficiency of a good but not linear set union algorithm", Journal of the ACM 22(2), 1975, the inverse Ackermann bound.
• Sedgewick and Wayne, Algorithms, section 1.5 (union-find).
• Skiena, The Algorithm Design Manual: the union-find section in the weighted graph chapter, where it speeds up Kruskal's algorithm, and the catalog entry on set data structures.
• Galler and Fischer, "An improved equivalence algorithm", Communications of the ACM 7(5), 1964, the original disjoint-set forest.
• The NeetCode roadmap and the Blind 75, Grind 169, NeetCode 150 and NeetCode 250 lists, used as the index of topics and problems.