Disjoint Set(Union Find木)を一から書けるようにする

将来のコーディング試験に備えて。そのうちfenwick treeとかほかのも書くかもしれない。(typedef long long ll;とかrepマクロの定義は省略する)

あくまで自分のための備忘録なので、初めて組む人はこっち見た方がいいと思う Union-Find Tree を理解する！素集合系を扱うデータ構造 | アルゴリズムロジック

コーディングの順番としては、 ノードの定義→findRoot→isSame→unite(→Setsize) とするのが自然な気がする。

まずは普通にノードを定義する。集合の大きさがよく問われる気がするので、sizeによる併合を考えておく。

private:
    vector<ll> sizes, parents;
public:
    disjoint_set(ll N) {
        sizes.resize(N);
        parents.resize(N);
        rep(i, N) {
            sizes[i] = 1;
            parents[i] = i;
        }
    }

sizesとか直で見れないようにしておく。後でわかるが、sizes[i]はノードiを含む集合の大きさになるとは限らないからだ（sizes[iの親ノード]が集合の大きさになる）。

次に集合の代表ノード(根)を見つける。

 ll findRoot(ll x) {
        if(parents[x] != x) parents[x] = findRoot(parents[x]); // 経路圧縮
        return parents[x];
    }

経路圧縮してO(log(N))にする。 そして、uniteとisSameを実装する。isSameからuniteが自然かな。uniteは「長いものに巻かれろ」という覚え方で、小さい方から大きい方に接続する。

 bool isSame(ll x, ll y) {
        return findRoot(x) == findRoot(y);
    }

    bool unite(ll x, ll y) {
        if (isSame(x, y)) return false;
        x = findRoot(x);
        y = findRoot(y);
        if (sizes[x] > sizes[y]) { // サイズが小さい方を大きい方につける
            parents[y] = x;
            sizes[x] += sizes[y];
        } else {
            parents[x] = y;
            sizes[y] += sizes[x];
        }
        return true;
    }

これだとunite時に無駄に二回findRootしていることになるが、まあ定数倍だし、こっちの方がわかりやすいのでこのままにしておこうと思う。小さい木を大きい方につけることで、経路圧縮と合わせて計算量はほとんど定数倍(O(α))に落ちる。 最後にSetsizeを実装して終わり。

 ll Setsize(ll x) {
        return sizes[findRoot(x)];
    }   

全体としてはこんな感じ：

struct disjoint_set{
private:
    vector<ll> sizes, parents;
public:
    disjoint_set(ll N) {
        sizes.resize(N);
        parents.resize(N);
        rep(i, N) {
            sizes[i] = 1;
            parents[i] = i;
        }
    }

    ll findRoot(ll x) {
        if(parents[x] != x) parents[x] = findRoot(parents[x]); // 経路圧縮
        return parents[x];
    }

    bool isSame(ll x, ll y) {
        return findRoot(x) == findRoot(y);
    }

    bool unite(ll x, ll y) {
        if (isSame(x, y)) return false;
        x = findRoot(x);
        y = findRoot(y);
        if (sizes[x] > sizes[y]) { // サイズが小さい方を大きい方につける
            parents[y] = x;
            sizes[x] += sizes[y];
        } else {
            parents[x] = y;
            sizes[y] += sizes[x];
        }
        return true;
    }

    ll Setsize(ll x) {
        return sizes[findRoot(x)];
    }
};

verifiedする：

judge.u-aizu.ac.jp

よさそう。