undoable union find sum

ソースコード

 1#include <vector>
 2using namespace std;
 3
 4// UndoableUnionFindSum
 5namespace titan23 {
 6
 7    template<typename T>
 8    class UndoableUnionFindSum {
 9      private:
10        int _n, _group_count;
11        T _e;
12        vector<int> _parents;
13        vector<T> _all_sum, _one_sum;
14        vector<tuple<int, int, T>> _history;
15
16      public:
17        UndoableUnionFindSum() {}
18        UndoableUnionFindSum(int n, T e) :
19            _n(n),
20            _group_count(n),
21            _e(e),
22            _parents(n, -1),
23            _all_sum(n, e),
24            _one_sum(n, e) {}
25
26        void undo() {
27            auto [y, py, all_sum_y] = _history.back();
28            _history.pop_back();
29            if (y == -1) return;
30            auto [x, px, all_sum_x] = _history.back();
31            _history.pop_back();
32            ++_group_count;
33            _parents[y] = py;
34            _parents[x] = px;
35            T s = (_all_sum[x] - all_sum_y - all_sum_x) / (-py-px) * (-py);
36            _all_sum[y] += s;
37            _all_sum[x] -= all_sum_y + s;
38            _one_sum[x] -= _one_sum[y];
39        }
40
41        int root(int x) const {
42            while (_parents[x] >= 0) {
43                x = _parents[x];
44            }
45            return x;
46        }
47
48        bool unite(int x, int y) {
49            x = root(x);
50            y = root(y);
51            if (x == y) {
52                _history.emplace_back(-1, -1, _e);
53                return false;
54            }
55            if (_parents[x] > _parents[y]) swap(x, y);
56            _group_count -= 1;
57            _history.emplace_back(x, _parents[x], _all_sum[x]);
58            _history.emplace_back(y, _parents[y], _all_sum[y]);
59            _all_sum[x] += _all_sum[y];
60            _one_sum[x] += _one_sum[y];
61            _parents[x] += _parents[y];
62            _parents[y] = x;
63            return true;
64        }
65
66        int size(const int x) const {
67            return -_parents[root(x)];
68        }
69
70        bool same(const int x, const int y) const {
71            return root(x) == root(y);
72        }
73
74        void add_point(int x, const T v) {
75            while (x >= 0) {
76                _one_sum[x] += v;
77                x = _parents[x];
78            }
79        }
80
81        void add_group(int x, const T v) {
82            x = root(x);
83            _all_sum[x] += v * size(x);
84        }
85
86        int group_count() const {
87            return _group_count;
88        }
89
90        T group_sum(int x) const {
91            x = root(x);
92            return _one_sum[x] + _all_sum[x];
93        }
94    };
95} // namespace titan23

仕様

Warning

doxygenfile: Cannot find file “titan_cpplib/data_structures/undoable_union_find_sum.cpp