Step 6 of 6

Challenge: union-find

Some problems are about groups that keep merging: which computers can reach each other as cables are added, which accounts belong to the same person as links are found. Re-running a BFS after every merge would be slow. Union-find (also called disjoint set union) answers "are these two in the same group?" and "merge these two groups" in nearly O(1) each.

Each group is stored as a tree through a parent array. The root of a tree identifies its group: find(x) follows parents up to the root. Two items are in the same group exactly when they have the same root, and merging two groups just points one root at the other.

#include <iostream>
#include <numeric>
#include <vector>

std::vector<int> parent;

int find(int x) {
    while (parent[x] != x) x = parent[x];     // walk up to the root
    return x;
}

void unite(int a, int b) {
    parent[find(a)] = find(b);               // hang one root under the other
}

int main() {
    parent.resize(6);
    std::iota(parent.begin(), parent.end(), 0);   // everyone is their own group: 0 1 2 3 4 5
    unite(0, 1);
    unite(2, 3);
    unite(1, 3);                                  // joins {0,1} with {2,3}
    std::cout << (find(0) == find(2)) << (find(0) == find(4)) << "\n";

    int groups = 0;
    for (int i = 0; i < 6; i++) groups += (find(i) == i);   // count roots
    std::cout << groups << " groups\n";
}

10
3 groups

The groups are {0, 1, 2, 3}, {4} and {5}.

Making it fast

This simple version can build long chains, making find O(n). Two tricks fix that:

• Union by size: keep a size array, and always attach the smaller tree under the root of the larger one. Trees then stay shallow (height at most log n).
• Path compression: during find, after locating the root, point every node you passed directly at it. The next find on any of them is a single step.

int find(int x) {
    int root = x;
    while (parent[root] != root) root = parent[root];
    while (parent[x] != root) {            // second pass: repoint the path
        int next = parent[x];
        parent[x] = root;
        x = next;
    }
    return root;
}

Together they make each operation effectively constant time. Keep a groups_ counter that starts at n and drops by one on every successful merge, and read a set's size from its root.

Your task

Write class UnionFind:

• explicit UnionFind(int n): n separate sets
• int find(int x) with path compression
• bool unite(int a, int b): merge; returns false if already in the same set
• int groups() const: number of separate sets
• int size_of(int x): size of x's set

Challenge 2

Write int count_components(int n, const std::vector<std::pair<int, int>>& edges) for an undirected graph using union-find: start with n groups and take one away for every edge that joins two different groups. Write a small find with path compression.

Challenge 3

Kruskal's algorithm finds a minimum spanning tree: sort the edges by weight, then take each edge whose ends are in different union-find groups. Write long long mst_weight(int n, std::vector<WEdge> edges), the total weight of the tree, or -1 if the graph isn't connected.

Previous: Dijkstra's shortest paths