UnionFind木を改造する話

　UnioinFind木は便利です。以下のようなかたちで実装できます。

```c++
struct UnionFind{
    vector<int> par;
    UnionFind(int n){
        par.resize(n,-1);
    }
    int root(int x){
        if(par[x]<0){
            return x;
        }
        return par[x] = root(par[x]); // ここで経路圧縮をするとうれしい
    }
    int size(int x){
        return -par[root(x)];
    }
    int merge(int x,int y){
        int rx = root(x);
        int ry = root(y);
        if(rx==ry){
            return rx;
        }
        if(size(rx)<size(ry)){
            swap(rx,ry); // ここで union by size するとうれしい
        }
        par[rx] += par[ry]; // ここで足しておくと実質的な size になってうれしい
        par[ry] = rx;
        return rx; // 実装の方針にもよるが、連結後の代表元が返ってくるとうれしい
    }
};
```

# エッジの個数を持たせる

https://atcoder.jp/contests/abc292/tasks/abc292_d

　連結成分ごとにエッジの個数を持たせるとこういう問題を解きやすくなります。

```c++
struct UnionFind{
    vector<int> par;
    vector<int> edge;
    UnionFind(int n){
        par.resize(n,-1);
        edge.resize(n,0);
    }
    int root(int x){
        if(par[x]<0){
            return x;
        }
        return par[x] = root(par[x]);
    }
    int size(int x){
        return -par[root(x)];
    }
    int merge(int x,int y){
        int rx = root(x);
        int ry = root(y);
        if(rx==ry){
            edge[rx]++; // 被っている場合もエッジだけは足す
            return rx;
        }
        if(size(rx)<size(ry)){
            swap(rx,ry);
        }
        par[rx] += par[ry];
        par[ry] = rx;
        edge[rx] += edge[ry]+1; // 被っていない場合も当然エッジは足す
        return rx;
    }
    int get_edge_cnt(int x){ // root を介する場合は、直接アクセスするよりも、こういう get 関数を用意した方が使いやすい
        return edge[root(x)];
    }
};
int main(){
    int n, m;
    cin >> n >> m;
    UnionFind uf(n);

    for (int i = 0; i < m; ++i) {
        int a, b;
        cin >> a >> b;
        a--; b--;
        uf.merge(a, b);
    }

    for (int i = 0; i < n; ++i) {
        if (uf.size(i) != uf.get_edge_cnt(i)) {
            cout << "No" << endl;
            return 0;
        }
    }

    cout << "Yes" << endl;
}
```

# 代表元を持たせる

　ルートとは別に、例えば「最小のノード」などを返すようにできます。

```c++
struct UnionFind {
    vector<int> par;
    vector<int> min_node; // 各集合の最小ノードを保持

    UnionFind(int n) {
        par.resize(n, -1);
        min_node.resize(n);
        for (int i = 0; i < n; i++) {
            min_node[i] = i;
        }
    }

    int root(int x) {
        if (par[x] < 0) {
            return x;
        }
        return par[x] = root(par[x]);
    }

    int size(int x) {
        return -par[root(x)];
    }

    int get_min_node(int x) {
        return min_node[root(x)];
    }

    int merge(int x, int y) {
        int rx = root(x);
        int ry = root(y);
        if (rx == ry) {
            return rx;
        }
        if (size(rx) < size(ry)) {
            swap(rx, ry);
        }
        min_node[rx] = min(get_min_node(rx), get_min_node(ry)); // 最小ノードを更新
        par[rx] += par[ry];
        par[ry] = rx;
        return rx;
    }
};
```

　要は `merge` のときにうまく都合の良い値を抽出できれば良いです。それには、作用させる順番を区別しないという性質が必要です。[セグ木の改造](https://qiita.com/hibit/items/0f5d058dee989437f7d4) のときでも似たようなことを話しました。つまり **モノイド** です。

## 最小代表限の飛び道具的実装

　実は、最小の元をルートにするようにすれば、代表元の情報は必要ないです。union by size が崩れますが、実は、

- union by size
- 経路圧縮
 
　は、片方だけでも $O(\log{N})$ になる（[参考記事](https://37zigen.com/union-find-complexity-1/)）ので、経路圧縮さえ生きていれば十分速いです。このことは、UnionFind にダイナミックな改造を施す場合に役に立つかもしれません。

```c++
struct UnionFind{
    vector<int> par;
    UnionFind(int n){
        par.resize(n,-1);
    }
    int root(int x){
        if(par[x]<0){
            return x;
        }
        return par[x] = root(par[x]);
    }
    int size(int x){
        return -par[root(x)];
    }
    int merge(int x,int y){
        int rx = root(x);
        int ry = root(y);
        if(rx==ry){
            return rx;
        }
        if(rx>ry){ // サイズではないので注意
            swap(rx,ry);
        }
        par[rx] += par[ry];
        par[ry] = rx;
        return rx;
    }
};
```

　このようにトリッキーな実装もできます。

　ちなみに、一次元の **連続した** 連結ノードを想定するなら、最小元＋サイズが自動的に最大となるので、実は最小元と最大元を付加情報として持たなくてもそれらを復元できます。

```c++
struct UnionFind{
    vector<int> par;
    UnionFind(int n){
        par.resize(n,-1);
    }
    int root(int x){
        if(par[x]<0){
            return x;
        }
        return par[x] = root(par[x]);
    }
    int size(int x){
        return -par[root(x)];
    }
    int merge(int x,int y){
        int rx = root(x);
        int ry = root(y);
        if(rx==ry){
            return rx;
        }
        if(rx>ry){ // サイズではないので注意
            swap(rx,ry);
        }
        par[rx] += par[ry];
        par[ry] = rx;
        return rx;
    }
    pair<int,int> get_minmax(int x){ // これで数直線上の [左端,右端) を得られる
        int rx = root(x);
        return {rx,rx+size(x)};
    }
};
```

　実際の問題での活用例。

 https://atcoder.jp/contests/abc370/submissions/66040012

# バケットの個数

　バケットの個数も付加情報として UnionFind に載せることができます。

```c++
struct UnionFind {
    vector<int> par;
    vector<map<int, int>> buc;
    
    UnionFind(int n) {
        par.resize(n, -1);
        buc.resize(n);
    }

    int root(int x) {
        if (par[x] < 0) return x;
        return par[x] = root(par[x]);
    }

    int size(int x) {
        return -par[root(x)];
    }

    void merge(int x, int y) {
        int rx = root(x);
        int ry = root(y);
        if (rx == ry) return;
        if (size(rx) < size(ry)) swap(rx, ry);
        par[rx] += par[ry];
        par[ry] = rx;
        for (auto [key, val] : buc[ry]) {
            buc[rx][key] += val;
        }
    }

    void init_bucket(const vector<int>& col) {
        int n = col.size();
        for (int i = 0; i < n; ++i) {
            buc[i][col[i]] = 1;
        }
    }
    
    map<int, int>& get_bucket(int x) { // これを参照扱いにしないと毎回コピーが発生してとんでもないことに
        return buc[root(x)];
    }
};
```

　実際の問題での活用例。

 https://atcoder.jp/contests/abc183/submissions/66008094