Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #include<bits/stdc++.h>
- using namespace std;
- struct StateDSU {
- int rootA, rootB, sizeA, sizeB;
- StateDSU(int _rootA, int _rootB, int _sizeA, int _sizeB) {
- rootA = _rootA;
- rootB = _rootB;
- sizeA = _sizeA;
- sizeB = _sizeB;
- }
- };
- class RollbackDSU {
- private:
- vector<int> par;
- vector<int> sz;
- vector<StateDSU> states;
- vector<int> snapshots;
- int setCount;
- public:
- // constructor
- RollbackDSU(int nodeCount) : setCount(nodeCount) {
- par.resize(nodeCount + 1);
- sz.assign(nodeCount + 1, 1);
- for(int i = 1; i <= nodeCount; ++i) {
- par[i] = i;
- }
- }
- private:
- // find the root of the component
- int getRoot(int node) {
- return node == par[node] ? node : node = getRoot(par[node]);
- }
- public:
- void addEdge(int a, int b) {
- int rootA = getRoot(a);
- int rootB = getRoot(b);
- int sizeA = sz[rootA];
- int sizeB = sz[rootB];
- states.push_back(StateDSU(rootA, rootB, sizeA, sizeB));
- if(rootA == rootB) {
- return;
- }
- setCount--;
- if(sizeA >= sizeB) {
- par[rootB] = rootA;
- sz[rootA] += sizeB;
- }
- else {
- par[rootA] = rootB;
- sz[rootB] += sizeA;
- }
- }
- bool rollBack() {
- if(states.empty()) {
- return false;
- }
- StateDSU lastState = states.back();
- states.pop_back();
- int rootA = lastState.rootA;
- int rootB = lastState.rootB;
- int sizeA = lastState.sizeA;
- int sizeB = lastState.sizeB;
- if(rootA == rootB) {
- return true;
- }
- setCount++;
- if(sizeA >= sizeB) {
- par[rootB] = rootB;
- sz[rootA] -= sizeB;
- }
- else {
- par[rootA] = rootA;
- sz[rootB] -= sizeA;
- }
- return true;
- }
- bool lastSnapshotRoll() {
- if(snapshots.empty()) {
- return false;
- }
- int lastSnapshot = snapshots.back();
- snapshots.pop_back();
- while(states.size() > lastSnapshot) {
- assert(rollBack());
- }
- return true;
- }
- void addSnapshot() {
- snapshots.push_back((int) states.size());
- }
- int getSetCount() {
- return setCount;
- }
- };
- int n, m;
- int main() {
- ios_base::sync_with_stdio(false);
- cin.tie(0);
- cin >> n >> m;
- RollbackDSU dsu(n);
- while(m--) {
- string op;
- cin >> op;
- if(op[0] == 'p')
- dsu.addSnapshot();
- else if(op[0] == 'r') {
- assert(dsu.lastSnapshotRoll());
- cout << dsu.getSetCount() << '\n';
- }
- else {
- int x, y;
- cin >> x >> y;
- dsu.addEdge(x, y);
- cout << dsu.getSetCount() << '\n';
- }
- }
- return 0;
- }
Advertisement
Add Comment
Please, Sign In to add comment