-
Notifications
You must be signed in to change notification settings - Fork 1
/
undo_dsu.hpp
53 lines (43 loc) · 1.03 KB
/
undo_dsu.hpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
#pragma once
#include "../template/template.hpp"
namespace lib {
struct undo_dsu {
public:
undo_dsu(int n) : n(n), data(n, -1) {}
int leader(int a) const {
if (data[a] < 0) return a;
return leader(data[a]);
}
int size(int a) const {
int x = leader(a);
return -data[x];
}
int merge(int a, int b) {
int x = leader(a);
int y = leader(b);
stack.push({x, data[x]});
stack.push({y, data[y]});
if (x == y) return x;
if (size(x) < size(y)) std::swap(x, y);
data[x] += data[y];
data[y] = x;
return x;
}
bool same(int a, int b) {
return leader(a) == leader(b);
}
void undo() {
assert(stack.size() >= 2);
auto [x, xval] = stack.top();
data[x] = xval;
stack.pop();
auto [y, yval] = stack.top();
data[y] = yval;
stack.pop();
}
private:
int n;
std::vector<int> data;
std::stack<std::pair<int, int>> stack;
};
} // namespace lib