☕ DSA Java Articles

Topics › Graphs › MinimumSpanningTree/Disjoint Set and Problems

Disjoint Set

Hard TUF ID 516 A2Z

Java source ↗ Edit this article ↗

Problem

Maintain n elements 0..n-1, initially each in its own group, under two operations:

The input is a sequence of operations [UNION, u, v] and [QUERY, u, v], and the output is the list of answers to the queries (“are u and v connected?”). This structure is called a disjoint set union (DSU) or union-find. It is the standard tool whenever a graph gains edges over time and you keep asking whether two vertices are connected. Rerunning BFS or DFS after every new edge would cost O(V + E) per question.

Example 1

n = 8
union(1,2) union(2,3) union(4,5) union(6,7) union(5,6)
query(3,7)  -> false     // groups: {1,2,3} {4,5,6,7}
union(3,7)
query(3,7)  -> true      // now one group {1,...,7}
query(1,4)  -> true

Example 2

n = 4, no unions
query(0,1) -> false, query(2,3) -> false     // every element is alone

All three approaches store each group as a rooted tree in a parent array. The root is the group’s representative and satisfies parent[root] == root. find walks up to the root, and union makes one root the child of the other. They differ only in how much effort they spend keeping the trees short.

The Java file defines each version as a nested class behind a small UnionFind interface, and one shared driver applies the operations.

interface UnionFind {
    int find(int x);
    void union(int a, int b);
}

static List<Boolean> run(UnionFind dsu, int[][] ops) {
    List<Boolean> answers = new ArrayList<>();
    for (int[] op : ops) {
        if (op[0] == UNION) dsu.union(op[1], op[2]);
        else answers.add(dsu.find(op[1]) == dsu.find(op[2]));
    }
    return answers;
}

Approach 1: Naive parent array

Intuition. The simplest version: find follows parent pointers until it reaches a root, and union finds both roots and points the first at the second. It is correct, because two elements share a root exactly when they were linked by some chain of unions. Nothing stops the trees from becoming long paths, though.

Algorithm.

  1. Start with parent[i] = i.
  2. find(x): while parent[x] != x, move to parent[x]. Return x.
  3. union(a, b): if find(a) != find(b), set parent[find(a)] = find(b).
static class NaiveDSU implements UnionFind {
    final int[] parent;

    NaiveDSU(int n) {
        parent = new int[n];
        for (int i = 0; i < n; i++) parent[i] = i;
    }

    public int find(int x) {
        while (parent[x] != x) x = parent[x];
        return x;
    }

    public void union(int a, int b) {
        int ra = find(a), rb = find(b);
        if (ra != rb) parent[ra] = rb;
    }
}

static List<Boolean> naive(int n, int[][] ops) {
    return run(new NaiveDSU(n), ops);
}

Complexity. find costs the depth of the tree, which can reach n - 1. The sequence union(0,1), union(1,2), …, union(n-2, n-1) builds exactly that chain, with 0 at the bottom, so each later find(0) walks the whole path. Over m operations this is O(m · n) in the worst case. Space O(n).

Approach 2: Union by rank with path compression

Intuition. Two independent ideas each fix a different cause of tall trees.

Union by rank. Keep a rank for every root, an upper bound on its tree’s height. When merging, hang the root with the smaller rank under the one with the larger rank. The height of the taller tree does not change. Only when the ranks are equal does the new root’s rank grow by one. A root of rank r therefore has at least 2^r elements in its tree (two rank r - 1 trees were needed to make it), so ranks, and heights, never exceed log2 n. The Java tests build a perfectly balanced sequence of merges on 1024 elements and check that the final rank is exactly 10.

Path compression. After find(x) discovers the root, point every node on the path from x directly at the root. The recursive version does it in one line: parent[x] = find(parent[x]). Later finds from any of those nodes take one step.

Together they give an amortised cost of O(α(n)) per operation, where α is the inverse Ackermann function. It is at most 4 for any n that fits in memory, so in practice every operation is constant time.

Algorithm.

  1. Start with parent[i] = i and rank[i] = 0.
  2. find(x): if x is a root return it, otherwise set parent[x] = find(parent[x]) and return that.
  3. union(a, b): find the roots ra and rb. If they are equal, stop. Otherwise attach the lower-rank root to the higher one; on a tie attach rb under ra and increment rank[ra].
static class DSUByRank implements UnionFind {
    final int[] parent, rank;

    DSUByRank(int n) {
        parent = new int[n];
        rank = new int[n];
        for (int i = 0; i < n; i++) parent[i] = i;
    }

    public int find(int x) {
        if (parent[x] == x) return x;
        return parent[x] = find(parent[x]);          // path compression: point straight at the root
    }

    public void union(int a, int b) {
        int ra = find(a), rb = find(b);
        if (ra == rb) return;
        if (rank[ra] < rank[rb]) {
            parent[ra] = rb;
        } else if (rank[ra] > rank[rb]) {
            parent[rb] = ra;
        } else {
            parent[rb] = ra;                         // equal ranks: the new root grows by one
            rank[ra]++;
        }
    }
}

static List<Boolean> byRank(int n, int[][] ops) {
    return run(new DSUByRank(n), ops);
}

Complexity. Amortised O(α(n)) per operation, so O(m · α(n)) for m operations. Space O(n) for two arrays.

Rank is only an upper bound once compression starts. Compression shortens paths without updating ranks, and that is fine: the bound is all the analysis needs.

Approach 3: Union by size with path compression

Intuition. Instead of rank, track size[root], the number of elements in the group, and hang the smaller group under the larger. The same argument applies: an element’s depth increases only when its group is merged into one at least as large, which at least doubles the size of its group. That can happen at most log2 n times. The complexity is identical to union by rank. The bonus is that size[find(x)] tells you how large x’s group is, which many problems need (largest island, accounts merge, counting pairs across components).

Algorithm.

  1. Start with parent[i] = i and size[i] = 1.
  2. find(x): the same path-compressing find.
  3. union(a, b): find both roots. If they differ, attach the root of the smaller group under the other and add its size to the new root.
static class DSUBySize implements UnionFind {
    final int[] parent, size;

    DSUBySize(int n) {
        parent = new int[n];
        size = new int[n];
        for (int i = 0; i < n; i++) {
            parent[i] = i;
            size[i] = 1;
        }
    }

    public int find(int x) {
        if (parent[x] == x) return x;
        return parent[x] = find(parent[x]);
    }

    public void union(int a, int b) {
        int ra = find(a), rb = find(b);
        if (ra == rb) return;
        if (size[ra] < size[rb]) {                   // hang the smaller tree under the larger one
            parent[ra] = rb;
            size[rb] += size[ra];
        } else {
            parent[rb] = ra;
            size[ra] += size[rb];
        }
    }
}

static List<Boolean> bySize(int n, int[][] ops) {
    return run(new DSUBySize(n), ops);
}

Dry run of Example 1 with union by size (n = 8):

operation roots compared result groups (root: size)
union(1,2) 1 (1) vs 2 (1) parent[2] = 1 1: 2
union(2,3) 1 (2) vs 3 (1) parent[3] = 1 1: 3
union(4,5) 4 (1) vs 5 (1) parent[5] = 4 1: 3, 4: 2
union(6,7) 6 (1) vs 7 (1) parent[7] = 6 1: 3, 4: 2, 6: 2
union(5,6) 4 (2) vs 6 (2) parent[6] = 4 1: 3, 4: 4
query(3,7) find(3) = 1, find(7) = 4 false  
union(3,7) 1 (3) vs 4 (4) parent[1] = 4 4: 7
query(3,7) find(3) = 4, find(7) = 4 true  

During the last find(3), the path 3 → 1 → 4 is compressed so that parent[3] = 4.

Complexity. Amortised O(α(n)) per operation; space O(n).

Edge cases and pitfalls