[c++] UnionFind

昔、Pythonで作成していたUnionFindのライブラリをC++で書き直しました。
グループ化された要素(木構造)の要素数にもとづいて、木の結合を行っています。

配列parentsで、木構造を表していますが、

vector<int> parents;

この配列の値に、要素数も保存するようにしています。
・値が0以上の時は、親のindexを意味しますが、
・値が負の値の時は、要素数を負の値にしたものが入力されています。
　　(例) index2の要素数5がとき、parents[2] = -5

[参考にしたサイト]
Union-Find木 [いかたこのたこつぼ]

code

#include <iostream>
#include <vector>
using namespace std;


// Debug用  配列の表示 ----------
template <class T>
void print_vector(vector<T> arr) {
    printf("[ ");
    for (int i = 0; i < arr.size(); i++) {
        cout << arr[i];
        if (i != (arr.size() - 1)) printf(", ");
    }
    printf(" ]\n");
}
// ------------------------------


// UnionFind
// 「木の要素数(size)」に基づいて、木の結合を行う
// 木の要素数は、配列parentsに負の値で保存(正の値の場合は、親の番号(index))
class UnionFind {
public:
    vector<int> parents;
    UnionFind(int n) : parents(n, -1) {}

    int find(int x) {
        if (parents[x] < 0) {
            return x;
        }
        parents[x] = find(parents[x]);
        return parents[x];
    }

    void unite(int x, int y) {
        int root_x = find(x);
        int root_y = find(y);
        if (root_x == root_y) {
            return;
        }

        int size_x = parents[root_x];
        int size_y = parents[root_y];
        if (size_x <= size_y) {
            parents[root_y] = root_x;
            parents[root_x] += size_y;
        } else {
            parents[root_x] = root_y;
            parents[root_y] += size_x;
        }
    }
};
// ------------------------------


int main() {
    UnionFind uf(10);
    printf("initialize  >> ");
    print_vector(uf.parents);

    vector<vector<int>> data = {{0, 4}, {6, 7}, {7, 8}};
    for (vector<int> d : data) {
        uf.unite(d[0], d[1]);
        printf("unite %d %d   >> ", d[0], d[1]);
        print_vector(uf.parents);
    }
}

[実行結果]
initialize >> [ -1, -1, -1, -1, -1, -1, -1, -1, -1, -1 ]
unite 0 4 >> [ -2, -1, -1, -1, 0, -1, -1, -1, -1, -1 ]
unite 6 7 >> [ -2, -1, -1, -1, 0, -1, -2, 6, -1, -1 ]
unite 7 8 >> [ -2, -1, -1, -1, 0, -1, -3, 6, 6, -1 ]

index0, 4が結合されていて、要素数2です。parents[0]に要素数が入力されており、値が-2になっています。

index6, 7, 8が結合されていて、要素数3です。parents[6]に要素数が入力されており、値が-3になっています。