Calculating Height and Size of Disjoint Set Union

The Disjoint Set Union (DSU), also known as the Union-Find data structure, is a fundamental tool in computer science used to manage a collection of non-overlapping sets. It is most commonly employed to solve connectivity problems, such as determining if two nodes in a graph belong to the same connected component or merging two distinct components.

While the basic functionality of DSU focuses on connectivity, advanced applications often require more granular information. Specifically, knowing the size of a set (the number of elements it contains) and its height (the depth of the tree structure) is crucial for analyzing subgraph scales, optimizing tree-based algorithms, and managing complex hierarchical data. This article provides a deep dive into how to efficiently maintain and calculate these two properties.

1. The Core Mechanics of DSU

Internally, a DSU is represented as a forest—a collection of disjoint trees. Each set is represented by a single tree, where the root node serves as the unique representative of that set.

The structure relies on two primary operations:

  • find(x): This operation traverses the tree from node x upwards until it reaches the root. The root is the identifier for the entire set.
  • union(x, y): This operation merges the sets containing x and y. If they are already in the same set, no action is taken. Otherwise, one root is made a child of the other, effectively merging the two trees.

Without optimization, a series of union operations could result in a tree that is essentially a long chain (a degenerate tree). In such a case, the find operation would take O(N) time, negating the efficiency of the data structure. To prevent this, we employ two critical heuristics:

  1. Path Compression: During a find operation, we make every node on the path point directly to the root. This flattens the tree significantly.
  2. Union by Rank or Size: When merging two trees, we attach the "smaller" tree to the root of the "larger" tree, preventing the height from growing unnecessarily.

2. Maintaining Set Size

The size of a set is defined as the total number of elements belonging to that specific component.

2.1 Implementation Logic

To track size, we maintain an auxiliary array, typically named sz[].

  • Initialization: Every element starts in its own set. Therefore, for every node i, we set sz[i] = 1.
  • Merging (Union by Size): When merging two sets with roots rx and ry, we compare their sizes. To keep the tree balanced, we attach the root of the smaller set to the root of the larger set. We then update the size of the new root:
    sz[new_root] = sz[new_root] + sz[old_root].
  • Querying: To find the size of the set containing element x, we first find its root using find(x) and then return sz[root].

Using Union by Size combined with Path Compression ensures that the amortized time complexity for all operations is O(α(N)), where $\alpha$ is the inverse Ackermann function—a value that remains effectively constant for all practical purposes.

3. Calculating and Managing Set Height

The height (or depth) of a set refers to the longest path from the root to any leaf node. In the context of DSU, we often use a concept called Rank.

3.1 Rank vs. Actual Height

It is important to distinguish between the actual height of the tree and its rank.

  • Rank is an upper bound on the height of the tree.
  • When using Path Compression, the actual height of the tree decreases, but updating the rank becomes computationally expensive. Therefore, we use the rank as a heuristic to guide the union operation rather than maintaining the exact depth.

3.2 Implementation Logic (Union by Rank)

To manage rank, we use a rank[] array.

  • Initialization: Each node starts with rank[i] = 0.
  • Merging:
    • If rank[rx] < rank[ry], make ry the parent of rx. The rank of ry remains unchanged.
    • If rank[rx] > rank[ry], make rx the parent of ry. The rank of rx remains unchanged.
    • If rank[rx] == rank[ry], pick one to be the parent (e.g., rx) and increment its rank: rank[rx]++.

This strategy ensures that the height of any tree grows logarithmically, specifically keeping the height bounded by O(log N) even without path compression.

3.3 Obtaining Real Height

If your specific application requires the exact depth of a node rather than an approximation, you must traverse the parent pointers manually:

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

Note: This manual traversal is generally discouraged within the core DSU logic as it can degrade performance.

4. Complete Implementation Examples

Below are optimized implementations in three widely used languages.

4.1 C++ (Union by Size + Path Compression)

#include <vector>
#include <numeric>
#include <algorithm>

class DSU {
public:
    DSU(int n) : parent(n), sz(n, 1) {
        std::iota(parent.begin(), parent.end(), 0);
    }

    int find(int x) {
        if (parent[x] == x) return x;
        return parent[x] = find(parent[x]); // Path Compression
    }

    void unite(int x, int y) {
        int rx = find(x);
        int ry = find(y);
        if (rx != ry) {
            // Union by Size
            if (sz[rx] < sz[ry]) std::swap(rx, ry);
            parent[ry] = rx;
            sz[rx] += sz[ry];
        }
    }

    int size(int x) {
        return sz[find(x)];
    }

private:
    std::vector<int> parent;
    std::vector<int> sz;
};

4.2 Java (Union by Rank + Path Compression)

public class DSU {
    private int[] parent;
    private int[] rank;

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

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

    public void union(int x, int y) {
        int rx = find(x);
        int ry = find(y);
        if (rx != ry) {
            // Union by Rank
            if (rank[rx] < rank[ry]) {
                parent[rx] = ry;
            } else if (rank[rx] > rank[ry]) {
                parent[ry] = rx;
            } else {
                parent[ry] = rx;
                rank[rx]++;
            }
        }
    }
}

4.3 Python (Union by Size + Path Compression)

class DSU:
    def __init__(self, n):
        self.parent = list(range(n))
        self.sz = [1] * n

    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])  # Path Compression
        return self.parent[x]

    def unite(self, x, y):
        rx, ry = self.find(x), self.find(y)
        if rx != ry:
            # Union by Size
            if self.sz[rx] < self.sz[ry]:
                rx, ry = ry, rx
            self.parent[ry] = rx
            self.sz[rx] += self.sz[ry]

    def get_size(self, x):
        return self.sz[self.find(x)]

5. Performance Analysis and Common Pitfalls

To ensure your DSU implementation is production-ready or competitive-programming grade, keep the following comparison in mind:

Feature Correct Approach Common Mistake
Path Compression Implement within the find function (recursive or iterative). Forgetting to update the parent, leading to $O(N)$ complexity.
Union Strategy Compare sz or rank to attach the smaller tree to the larger one. Arbitrarily picking a parent, which can lead to a degenerate tree.
Size Maintenance Update the sz value only at the root node during a union. Attempting to update sizes for all nodes, which is inefficient.
Complexity Goal Aim for amortized $O(\alpha(N))$. Relying on a simple union without any heuristic.

6. Summary

  • Set Size: Managed via a sz[] array. During a union, the size of the new root is the sum of the sizes of the two merged roots. This is highly effective for problems requiring component counting.
  • Set Height: Managed via a rank[] array. This provides a heuristic upper bound on the tree depth, ensuring logarithmic height growth.
  • Optimization Synergy: The combination of Path Compression and Union by Size/Rank is the "golden standard" for DSU, providing near-constant time complexity for all operations.

By mastering these two properties, you can extend the utility of the Disjoint Set Union from a simple connectivity checker to a powerful tool for complex structural analysis in graph theory and distributed systems.