Union-Find Size Optimization

Following the approach from the previous section, we perform the union(4,9) operation on the depicted union-find structure as shown below.

After the merge operation, the structure becomes:

It can be observed that the tree of this structure has relatively high levels. If the number of elements increases, the resulting overhead will be relatively large. Solving this issue is straightforward: during the specific pointing operation, we first check and point the root node of the smaller set to the root node of the larger set, which increases the likelihood of generating a tree with fewer levels.

When constructing the union-find set, an additional parameter, the sz array, is needed. sz[i] represents the number of elements in the set rooted at i.

// Constructor
public UnionFind3(int count){
    parent = new int[count];
    sz = new int[count];
    this.count = count;
    // Initialization, each parent[i] points to itself, indicating each element forms its own set
    for( int i = 0 ; i < count ; i ++ ){
        parent[i] = i;
        sz[i] = 1;
    }
}

During the merge operation, the direction of merging is determined based on the number of elements in the trees of the two elements.

public void unionElements(int p, int q){
    int pRoot = find(p);
    int qRoot = find(q);
    if( pRoot == qRoot )
        return;
    if( sz[pRoot] < sz[qRoot] ){
        parent[pRoot] = qRoot;
        sz[qRoot] += sz[pRoot];
    }
    else{
        parent[qRoot] = pRoot;
        sz[pRoot] += sz[qRoot];
    }
}

After optimization, the merge result is as follows, with 9 pointing to parent node 8.

Java Example Code

Source Code Download: Download

UnionFind3.java File Code:

package tutorialpro.union;

/**
 * Union-Find size optimization
 */
public class UnionFind3 {
    // parent[i] indicates the parent node of the first element
    private int[] parent;
    // sz[i] indicates the number of elements in the set rooted at i
    private int[] sz;
    // Number of elements
    private int count;

    // Constructor
    public UnionFind3(int count){
        parent = new int[count];
        sz = new int[count];
        this.count = count;
        // Initialization, each parent[i] points to itself, indicating each element forms its own set
        for( int i = 0 ; i < count ; i ++ ){
            parent[i] = i;
            sz[i] = 1;
        }
    }

    // Find process, find the set number corresponding to element p
    // O(h) complexity, h is the height of the tree
    private int find(int p){
        assert( p >= 0 && p < count );
        // Continuously query the parent node until reaching the root node
        // Root node characteristic: parent[p] == p
        while( p != parent[p] )
            p = parent[p];
        return p;
    }

    // Check if elements p and q belong to the same set
    // O(h) complexity, h is the height of the tree
    public boolean isConnected( int p , int q ){
        return find(p) == find(q);
    }

    // Merge the sets of elements p and q
    // O(h) complexity, h is the height of the tree
    public void unionElements(int p, int q){
        int pRoot = find(p);
        int qRoot = find(q);
        if( pRoot == qRoot )
            return;
        if( sz[pRoot] < sz[qRoot] ){
            parent[pRoot] = qRoot;
            sz[qRoot] += sz[pRoot];
        }
        else{
            parent[qRoot] = pRoot;
            sz[pRoot] += sz[qRoot];
        }
    }
}

int qRoot = find(q);
if (pRoot == qRoot)
    return;
// Based on the different sizes of the trees containing the two elements, determine the merge direction
// Merge the smaller set into the larger set
if (sz[pRoot] < sz[qRoot]) {
    parent[pRoot] = qRoot;
    sz[qRoot] += sz[pRoot];
} else {
    parent[qRoot] = pRoot;
    sz[pRoot] += sz[qRoot];
}