Skip to content

Union find add

c++

 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
struct UnionFind{
    vector<int> par; // par[i]:iの親の番号
    vector<int> s; // s[i]:iの属する木のノード数
    int n_group; // 木の数
    VI val;

    UnionFind(int N) : par(N) {
        for(int i=0; i<N; i++) par[i] = i;
        s.resize(N, 1);
        n_group = N;
        val.resize(N,0);
    }

    int root(int x){ // xが属する木の根
        if(par[x] == x) return x;
        int p = par[x];
        par[x] = root(par[x]);
        if(p != par[x]) val[x] += val[p];
        return par[x];
    }

    void unite(int x, int y){ // xとyの木を併合
        int rx = root(x);
        int ry = root(y);
        if(rx == ry) return;
        if(s[rx] > s[ry]) swap(rx, ry);
        par[rx] = ry;
        val[rx] -= val[ry];
        s[ry] += s[rx];
        s[rx] = 0;
        n_group--;
    }

    bool same(int x, int y){ // xとyが同じ木に属するかどうか
        return root(x) == root(y);
    }

    int size(int x){ // xが属する木のノード数
        return s[root(x)];
    }

    void add(int pos, int x){
        int r = root(pos);
        val[r] += x;
    }

    int query(int pos){
        int r = root(pos);
        return r == pos ? val[pos] : val[r] + val[pos];
    }
};