#include "collapse.h"
#include <bits/stdc++.h>
using namespace std;
class disjoint_set {
public:
disjoint_set() {}
disjoint_set(int n) : n_(n), comp_count(n), comp(n, -1) {}
int find_root(int x) {
return (comp[x] < 0) ? x : find_root(comp[x]);
}
int merge(int x, int y) {
x = find_root(x);
y = find_root(y);
if (x == y) {
return 0;
}
if (-comp[x] > -comp[y]) {
swap(x, y);
}
history.push_back({x, comp[x]});
history.push_back({y, comp[y]});
comp[y] += comp[x];
comp[x] = y;
comp_count--;
return 1;
}
void rollback(int x) {
comp_count += x;
for (int i = 0; i < 2 * x; i++) {
comp[history.back().first] = history.back().second;
history.pop_back();
}
}
void make_set() {
n_ += 1;
comp.emplace_back(-1);
}
int size(int x) {
return -comp[find_root(x)];
}
int size() {
return comp_count;
}
private:
int n_;
int comp_count;
vector<int> comp;
vector<pair<int, int>> history;
};
vector<int> simulateCollapse(int N, vector<int> T,
vector<int> X,
vector<int> Y,
vector<int> W,
vector<int> P) {
int C = (int) T.size();
int Q = (int) W.size();
vector<int> ans(Q);
vector<array<int, 4>> events;
for (int i = 0; i < C; i++) {
if (X[i] > Y[i]) {
swap(X[i], Y[i]);
}
if (T[i] == 0) {
events.push_back({i, -1, X[i], Y[i]});
} else {
events.push_back({i, -2, X[i], Y[i]});
}
}
for (int i = 0; i < Q; i++) {
events.push_back({W[i], i, P[i], P[i] + 1});
}
for (int rep = 0; rep < 2; rep++) {
struct edge {
int u, v;
edge() : u(0), v(0) {}
edge(int u, int v) : u(u), v(v) {}
const bool operator < (const edge &o) const {
return make_pair(u, v) < make_pair(o.u, o.v);
}
};
struct interval {
int l, r;
interval() : l(0), r(0) {}
interval(int l, int r) : l(l), r(r) {}
};
const int BLOCK = 777;
const int INF = 1e8;
set<edge> active_edges;
sort(begin(events), end(events));
for (int current = 0; current < (int) events.size(); current += BLOCK) {
set<edge> changed_edges;
for (int i = current; i < min(int(events.size()), current + BLOCK); i++) {
if (events[i][1] < 0) {
changed_edges.insert({events[i][2], events[i][3]});
}
}
map<edge, int> active_time;
vector<edge> edges_permanent;
for (auto e : active_edges) {
if (changed_edges.count(e) == 0) {
edges_permanent.push_back(e);
} else {
active_time[e] = 0;
}
}
vector<array<int, 3>> queries; // (x-coord, time, id)
vector<pair<edge, interval>> edges_temp;
for (int i = current; i < min(int(events.size()), current + BLOCK); i++) {
if (events[i][1] == -1) {
active_time[edge(events[i][2], events[i][3])] = events[i][0];
} else if (events[i][1] == -2) {
edge e = edge(events[i][2], events[i][3]);
edges_temp.push_back(make_pair(e, interval(active_time[e], events[i][0])));
active_time.erase(e);
} else {
queries.push_back({events[i][2], events[i][0], events[i][1]});
}
}
for (auto e : active_time) {
edges_temp.push_back(make_pair(e.first, interval(e.second, INF)));
}
sort(begin(edges_permanent), end(edges_permanent), [](const edge &e1, const edge &e2) {
return e1.v < e2.v;
});
reverse(begin(edges_permanent), end(edges_permanent));
sort(begin(queries), end(queries));
disjoint_set dsu(N);
for (auto q : queries) {
while (!edges_permanent.empty() && edges_permanent.back().v <= q[0]) {
int u = edges_permanent.back().u;
int v = edges_permanent.back().v;
edges_permanent.pop_back();
dsu.merge(u, v);
}
int rollbacks = 0;
for (auto e : edges_temp) {
if (e.first.v <= q[0] && e.second.l <= q[1] && q[1] < e.second.r) {
rollbacks += dsu.merge(e.first.u, e.first.v);
}
}
ans[q[2]] += dsu.size() - (N - q[0] - 1);
dsu.rollback(rollbacks);
}
for (int i = current; i < min(int(events.size()), current + BLOCK); i++) {
if (events[i][1] == -1) {
assert(active_edges.count({events[i][2], events[i][3]}) == 0);
active_edges.insert({events[i][2], events[i][3]});
} else if (events[i][1] == -2) {
assert(active_edges.count({events[i][2], events[i][3]}) == 1);
active_edges.erase({events[i][2], events[i][3]});
}
}
}
for (int i = 0; i < (int) events.size(); i++) {
events[i][2] = N - events[i][2] - 1;
events[i][3] = N - events[i][3] - 1;
swap(events[i][2], events[i][3]);
}
}
return ans;
}
