#include "collapse.h"
#include <bits/stdc++.h>
using namespace std;
constexpr int INF = 1e9;
typedef pair<int, int> pii;
struct weightedDSU {
    vector<int> par, weight, index;
    vector<int> st;
    int N;
    vector<pii> rpar, rweight, rst;
    vector<array<size_t, 3>> reverse;
    weightedDSU(int n) : par(n), weight(n, INF), index(n) {
        iota(par.begin(), par.end(), 0);
        iota(index.begin(), index.end(), 0);
        std::random_device rd;
        std::mt19937 gen(rd());
        ranges::shuffle(index, gen);
        N = 1 << 32 - __builtin_clz(1e6);
        st.assign(2 * N, 0);
    }
    int getRoot(int v, int w = INF - 1) {
        while (weight[v] <= w) {
            while (weight[v] >= weight[par[v]]) {
                rpar.emplace_back(v, par[v]);
                par[v] = par[par[v]];
            }
            v = par[v];
        }
        return v;
    }
    int mainEdge(int u, int v) {
        if (getRoot(u) != getRoot(v))
            return -1;
        while (par[u] != v && par[v] != u) {
            if (weight[u] < weight[v])
                u = par[u];
            else
                v = par[v];
        }
        if (par[u] == v)
            return u;
        return v;
    }
    void connect(int u, int v, int w) {
        while (u != v) {
            u = getRoot(u);
            v = getRoot(v);
            if (index[u] < index[v])
                swap(u, v);
            int tpar = par[u], tw = weight[u];
            rpar.emplace_back(u, par[u]);
            par[u] = v;
            rweight.emplace_back(u, weight[u]);
            weight[u] = w;
            v = tpar;
            w = tw;
        }
    }
    void addEdge(int u, int v, int w) {
        reverse.push_back({rpar.size(), rweight.size(), rst.size()});
        int p = mainEdge(u, v);
        if (p == -1) {
            connect(u, v, w);
            rst.emplace_back(N + w, -1);
            for (int i = N + w; i > 0; i >>= 1)
                st[i]++;
        }
        else if (weight[p] > w) {
            rpar.emplace_back(p, weight[p]);
            par[p] = p;
            rst.emplace_back(N + weight[p], 1);
            for (int i = N + weight[p]; i > 0; i >>= 1)
                st[i]--;
            rweight.emplace_back(p, weight[p]);
            weight[p] = INF;
            rst.emplace_back(N + w, -1);
            for (int i = N + w; i > 0; i >>= 1)
                st[i]++;
            connect(u, v, w);
        }
    }
    int query(int w) {
        int ans = 0;
        for (int l = N, r = N + w + 1; l < r; l >>= 1, r >>= 1) {
            if (l % 2) ans += st[l++];
            if (r % 2) ans += st[--r];
        }
        return ans;
    }
    void rollback() {
        auto [rpars, rws, rsts] = reverse.back();
        reverse.pop_back();
        while (rpar.size() > rpars) {
            auto [v, pv] = rpar.back();
            rpar.pop_back();
            par[v] = pv;
        }
        while (rweight.size() > rws) {
            auto [v, wv] = rweight.back();
            rweight.pop_back();
            weight[v] = wv;
        }
        while (rst.size() > rsts) {
            auto [v, add] = rst.back();
            rst.pop_back();
            for (; v > 0; v >>= 1)
                st[v] += add;
        }
    }
};
struct edge {
    int s, e, u, v;
};
void DandC(weightedDSU& dsu, int tl, int tr, const vector<edge>& edges, const vector<vector<pii>>& query, vector<int>& ans) {
    int rc = 0;
    int tm = (tl + tr) / 2;
    vector<edge> edgel, edger;
    for (auto [s, e, u, v] : edges) {
        if (s == tl && e == tr) {
            rc++;
            dsu.addEdge(u, v, max(u, v));
        }
        else {
            if (s < tm)
                edgel.push_back({s, min(tm, e), u, v});
            if (e > tm)
                edger.push_back({max(s, tm), e, u, v});
        }
    }
    if (tl + 1 == tr) {
        for (auto [t, w] : query[tl])
            ans[t] -= dsu.query(w);
    }
    else {
        DandC(dsu, tl, tm, edgel, query, ans);
        DandC(dsu, tm, tr, edger, query, ans);
    }
    while (rc--)
        dsu.rollback();
}
vector<int> simulateCollapse(
    int N,
    vector<int> T,
    vector<int> X,
    vector<int> Y,
    vector<int> W,
    vector<int> P
) {
    int Q = W.size();
    int M = T.size();
    vector<int> ans(Q, N);
    vector<vector<pii>> queryf(M), queryr(M);
    for (int i = 0; i < Q; i++) {
        queryf[W[i]].emplace_back(i, P[i]);
        queryr[W[i]].emplace_back(i, N - 2 - P[i]);
    }
    map<pii, int> mp;
    vector<edge> edges;
    for (int i = 0; i < M; i++) {
        if (X[i] > Y[i])
            swap(X[i], Y[i]);
        if (T[i] == 0)
            mp[{X[i], Y[i]}] = i;
        else {
            edges.push_back({mp[{X[i],Y[i]}], i, X[i], Y[i]});
            mp.erase({X[i], Y[i]});
        }
    }
    for (auto [p, s] : mp)
        edges.push_back({s, M, p.first, p.second});
    weightedDSU dsu(N);
    DandC(dsu, 0, M, edges, queryf, ans);
    for (auto& [s, e, u, v] : edges)
        u = N - 1 - u, v = N - 1 - v;
    DandC(dsu, 0, M, edges, queryr, ans);
    return ans;
}
| # | Verdict  | Execution time | Memory | Grader output | 
|---|
| Fetching results... | 
| # | Verdict  | Execution time | Memory | Grader output | 
|---|
| Fetching results... | 
| # | Verdict  | Execution time | Memory | Grader output | 
|---|
| Fetching results... | 
| # | Verdict  | Execution time | Memory | Grader output | 
|---|
| Fetching results... |