weight union find

ソースコード

 1#include <vector>
 2using namespace std;
 3
 4namespace titan23 {
 5    template<typename T>
 6    class WeightedUnionFind {
 7      private:
 8        int n, group_members;
 9        vector<int> par;
10        vector<T> weight;
11
12      public:
13        WeightedUnionFind() {}
14        WeightedUnionFind(int n) : n(n), group_members(n), par(n, -1), weight(n, 0) {}
15
16        int root(int x) {
17            vector<int> path = {x};
18            while (par[x] >= 0) {
19                x = par[x];
20                path.emplace_back(x);
21            }
22            int a = path.back();
23            path.pop_back();
24            while (!path.empty()) {
25                x = path.back();
26                path.pop_back();
27                weight[x] += weight[par[x]];
28                par[x] = a;
29            }
30            return a;
31        }
32
33        bool unite(int x, int y, T w) {
34            int rx = this->root(x);
35            int ry = this->root(y);
36            if (rx == ry) {
37                return this->diff(x, y) == w ? true : false;
38            }
39            w += weight[x] - weight[y];
40            group_members--;
41            if (par[rx] > par[ry]) {
42                swap(rx, ry);
43                w = -w;
44            }
45            par[rx] += par[ry];
46            par[ry] = rx;
47            weight[ry] = w;
48            return true;
49        }
50
51        int size(int x) {
52            return -par[root(x)];
53        }
54
55        bool same(int x, int y) {
56            return root(x) == root(y);
57        }
58
59        int group_count() const {
60            return group_members;
61        }
62
63        T diff(int x, int y) {
64            assert(same(x, y));
65            return weight[y] - weight[x];
66        }
67    };
68}

仕様

Warning

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