Zenn
🟢

UnionFindを実装する

2023/06/06に公開
競技プログラミング
tech

競技プログラミングに取り組んでいて，ある2頂点を結ぶパスが存在するかどうかを判定する問題(e.g. E - Good Graph)に出会う機会が増えてきました．
いままでは出来合いのUnionFindクラスを利用してきたのですが，友人から「競プロをやる上で，中身がブラックボックスのものを使うのではなく，どう動いているかを理解していたほうがいい」というようなアドバイスをもらい，自前で実装することにしました．

UnionFindとは

UnionFindとは，与えられたN頂点に対して，それらをいくつかの集合に分けるデータ構造のことを言います．集合間の結合(Union)とどの集合に属するか(FInd)の操作を高速に行うことができます．

プログラミング的には，N個の要素を持つ配列で表現できます．i番目の要素の親がA_iであると考え，根には-(その集合に含まれる要素の数)を保存しA_i < 0であるような場合その要素が根であると考えます．
ある頂点u, vについて，それらの親をたどって見つかる根が同じ値となるかを調べることによって， その2頂点が同じ集合に属するか否かを判定することができます．

コンストラクタ

はじめ，すべての頂点は独立している，すなわち根に当たります．独立した要素の要素数はもちろん1なので，すべて-1で初期化します．

UnionFind(int size) : uf(size) {
	for (int i = 0; i < size; i++) uf[i] = -1;
}

root

rootは，ある頂点nの親をたどっていき，その集合の根を返します．
その際経路圧縮を行い，レベル1の木となるようにすることで実行の高速化を図っています．

int root(int n) {
    if (uf[n] >= 0) uf[n] = root(uf[n]);
    return n;
}

connected

与えられた2頂点a, bが同じ集合に属する場合はa, そうでない場合はbを返します．

bool connected(int a, int b) { return root(a) == root(b); }

marge

２つの木を併合する操作です．与えられた2頂点a, bがすでに同じ集合に属する場合は操作を行わないようにしています．
また，この操作の際に集合のサイズを更新します．

void marge(int a, int b) {
    int root_a = root(a);
    int root_b = root(b);
    if (root_a != root_b) {
        if (root_a > root_b) swap(root_a, root_b);
        uf[root_a] += root_b;
        uf[root_b] = root_a;
    }
}

size

与えられた要素nについて，nが属する集合にいくつの要素が含まれるかを調べて返します．
ルートノードは-(要素数)なので，これにより簡潔に書くことができます．

int size(int n) { return -uf[root(n)]; }

全体

class UnionFind {
   private:
    vector<int> uf;

   public:
    UnionFind(int n) : uf(n) {
        for (int i = 0; i < n; i++) uf[i] = -1;
    }

    int root(int n) {
        if (uf[n] >= 0) uf[n] = root(uf[n]);
        return n;
    }

    bool connected(int a, int b) { return root(a) == root(b); }

    void marge(int a, int b) {
        int root_a = root(a);
        int root_b = root(b);
        if (root_a != root_b) {
            if (root_a > root_b) swap(root_a, root_b);
            uf[root_a] += root_b;
            uf[root_b] = root_a;
        }
    }

    int size(int n) { return -uf[root(n)]; }
};

最後に

このコードではじめに載せた問題を問題なく解くことができました．
他の人の提出を見ていると自分のコードの半分以下の時間で解いているような提出がいくつか見られたので，全体的にさらなる高速化ができるのでは？と思っています．このあたりの改善ができたらまた記事を書くか，本記事に追記したいと思います．

(6/6 18:20 : sizeの改善を行いました)

Discussion