Disjoint Set Union

#include<bits/stdc++.h>
#define ll long long

struct DSU // disjoint-set union
{
// DSU NOTE:
//      1. vertex can be 0- or 1-indexed, both possible.
//      2. by default, use get_parent_simple ( O(logN)). Change to ~O(1) by replace all get_parent_simple => get_parent_compress
    int *parent;
    int *ssize;
    
    DSU(int n) {
        parent = new int[n+2];
        ssize  = new int[n+2];
        
        for (int i = 0; i <= n; ++i) {
            parent[i] = i;
            ssize[i]  = 1;
        }
    
    }

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

    int get_parent_compress(int x) {
        
        if (x == parent[x])
            return x;
            
        parent[x] = get_parent_simple(parent[x]);
        return parent[x];
    }

    void join_set(int x, int y) {
        int pa_x = get_parent_simple(x);
        int pa_y = get_parent_simple(y);
        
        if (pa_x != pa_y) {
            if (ssize[pa_x] > ssize[pa_y]) {
                swap(pa_x, pa_y);
            }
            parent[pa_x] = pa_y;
            ssize[pa_y] += ssize[pa_x];
        }
    }

};

Leave a Reply