DSU with rollback

dsuRollBack.cpp

#include <bits/stdc++.h>

using namespace std;
typedef long long ll;

class DSU {
    vector<ll> p;             // parent if positive, size if negative, parent of root is the size of component
    stack<pair<ll, ll>> log;  // keep track of p[i] value history, first is index, second is value
    stack<ll> size;

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

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

        return root(p[x]);  // no path compression (reduce changes to p array)
    }

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

    void unite(ll a, ll b) {
        a = root(a), b = root(b);
        size.push(size.top() + 1);

        if (a == b) {
            log.emplace(0, -1);
            log.emplace(0, -1);
            size.top()--;

            return;
        }

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

        size.top() -= (p[a] != -1) + (p[b] != -1);

        log.emplace(a, p[a]);
        log.emplace(b, p[b]);  // remember to rollback 2 changes on rollback()

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

    ll cnt() {
        return size.top();
    }

    void rollback() {  // rollback 2 changes on p array
        for (ll i = 0; i < 2; i++) {
            auto [idx, val] = log.top();
            p[idx] = val;
            log.pop();
        }

        size.pop();
    }
};

int main() {
    DSU dsu(5);

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

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

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