Rollback Union Find with Mo's algorithm

ufrollbackmo.cpp

#include <bits/stdc++.h>
using namespace std;

#define rep(i, a, b) for(int i = a; i < (b); ++i)
#define trav(it, v) for(auto& it : v)
#define all(x) begin(x), end(x)
#define sz(x) (int)(x).size()
typedef long long ll;
typedef pair<int, int> pii;
typedef vector<int> vi;

struct UF {
    struct Merge {
        int u, v, parv;
        bool merged;
        Merge(int u, int v, int parv, bool merged) : u(u), v(v), parv(parv), merged(merged) {}
    };

    int merged = 0;
    vi par;
    int cmp;
    vector<Merge> ops;

    UF(int n) : par(n, -1) {}

    int find(int v) {
        if (par[v] < 0) return v;
        return find(par[v]);
    }

    void join(int u, int v) {
        u = find(u);
        v = find(v);
        if (u == v) {
            ops.emplace_back(-1, -1, -1, false);
            return;
        }
        merged++;
        if (-par[u] < -par[v]) swap(u, v);
        ops.emplace_back(u, v, par[v], true);
        par[u] += par[v];
        par[v] = u;
    }

    void rollback() {
        Merge m = ops.back();
        ops.pop_back();
        if (m.merged) {
            par[m.v] = m.parv;
            par[m.u] -= m.parv;
            merged--;
        }
    }
};

int add(int v, int l, int r, UF& uf, vector<vi>& F) {
    int res = 0;
    trav(it, F[v]) {
        if (l <= it && it < r) {
            uf.join(it, v);
            res++;
        }
    }
    return res;
}

struct Query {
    int name, l, r;

    bool operator<(const Query& ot) const {
        return r < ot.r;
    }
};

int main() {
	cin.tie(0)->sync_with_stdio(0);
	cin.exceptions(cin.failbit);

    int N, K;
    cin >> N >> K;

    int M;
    cin >> M;
    vector<vi> F(N);
    rep(i,0,M) {
        int a, b;
        cin >> a >> b;
        --a; --b;
        F[a].push_back(b);
        F[b].push_back(a);
    }

    int sq = (int)ceil(sqrt(N));
    int blocks = (N + sq - 1) / sq;
    vector<vector<Query>> Qs(blocks);
    int Q;
    cin >> Q;
    rep(i,0,Q) {
        int a, b;
        cin >> a >> b;
        --a; --b;
        Qs[a / sq].push_back({i, a, b + 1});
    }
    for (auto& v : Qs) sort(all(v));
    vi res(Q);

    rep(i,0,blocks) {
        if (Qs[i].empty()) continue;

        UF uf(N);
        int q = 0;
        while (q < sz(Qs[i]) && Qs[i][q].r <= (i + 1) * sq) {
            auto qr = Qs[i][q];
            int rb = 0;
            rep(v,qr.l,qr.r) {
                rb += add(v, qr.l, v, uf, F);
            }
            res[qr.name] = (qr.r - qr.l - uf.merged);
            rep(r,0,rb)  uf.rollback();
            q++;
        }
        int right = (i + 1) * sq;
        int left = (i + 1) * sq;
        for (; q < sz(Qs[i]); ++q) {
            auto qr = Qs[i][q];
            while (right < qr.r) {
                add(right, left, right, uf, F);
                right++;
            }
            int rb = 0;
            for (int v = left - 1; v >= qr.l; --v) {
                rb += add(v, v, right, uf, F);
            }
            res[qr.name] = (qr.r - qr.l - uf.merged);
            rep(r,0,rb) uf.rollback();
        }
    }
    trav(it, res) cout << it << endl;
}