Disjoint Set Union (DSU) C++ implementation

dsu.cpp

#include <bits/stdc++.h>

using namespace std;
typedef long long ll;

class DSU {
   private:
    vector<ll> p;  // parent if positive, size if negative, parent of root is the size of component

   public:
    DSU(ll N) {
        p = vector<ll>(N, -1);
    }

    ll root(ll x) {
        if (p[x] < 0) {
            return x;
        }

        return p[x] = root(p[x]); // path compression (keep tree height under 2)
    }

    bool sameSet(ll a, ll b) {
        return root(a) == root(b);
    }

    void unite(ll a, ll b) {
        a = root(a), b = root(b);

        if (a == b) {
            return;
        }

        if (p[a] > p[b]) { // small to large merging
            swap(a, b);
        }

        p[a] += p[b];
        p[b] = a;

        return;
    }

    ll size(ll x) {
        return -p[root(x)];
    }
};

int main() {
    DSU dsu(10);

    dsu.unite(1, 2);

    cout << dsu.sameSet(1, 2) << '\n';  // true
    cout << dsu.sameSet(1, 3) << '\n';  // false
    cout << dsu.size(1) << '\n';        // 2

    dsu.unite(1, 3);

    cout << dsu.sameSet(1, 2) << '\n';  // true
    cout << dsu.sameSet(1, 3) << '\n';  // true
    cout << dsu.size(1) << '\n';        // 3
}