Mister雑記

競プロやります。

部分永続Union-Findの実装

初めに

この記事では部分永続Union-Find(以下「部分永続UF」)の実装を解説するものであり、踏み込んだ原理の解説はしません。あくまでも「どうやって実装するのか?」という部分に重点を置いた記事になります。

しかしこの記事を読むにはまず部分永続UFの原理がある程度理解できていないと厳しいと思われます。ということで、私が参考にしたサイトを以下に記しておきます。

ちなみにUnion-Find(以下「UF」)の知識は必須なので、そちらを知らない方にはまだ本記事は早いかと思われます。ごめんね。

部分永続Union-Findとは?

そもそも部分永続UFとは、UFから多少の計算量を犠牲にして以下のクエリに対応できるようにしたものです。

  • 2つの頂点を含む木を結合する
  • t回結合を行った時点で、2つの頂点が同じ木に属するか(つまり連結か)調べる
  • t回uniteを行った時点での、ある頂点を含む木の頂点数を返す

UFとの違いは下線部のみ。つまり過去の状態も保持したUFとも言えます。それでいて計算量も空間量もUFと大差ないというチート並の性能を持つデータ構造なのです。

ただし1つ大きくUFに劣る点として、その性質上経路圧縮が行えないことが挙げられます。これにより、探索のならし計算量は O(\alpha(|V|))から O(\log |V|)に落ちます。まぁUFが異常なだけで十分高速なのですが。

実装

基本的にはUFに手を加えるだけでできるので、今回は以下の私のライブラリをイジることにします。

// 頂点数
const int V_NUM = 100000;

class UnionFind {
public:
    int par[V_NUM];  // 各頂点の親
    int rank[V_NUM]; // その頂点を根とする木の深さ
    int num[V_NUM];  // その頂点を根とする木の頂点数
    
    explicit UnionFind() {
        for (int i = 0; i < V_NUM; ++i) {
            par[i] = i;
        }
        fill(rank, rank + V_NUM, 0);
        fill(num, num + V_NUM, 1);
    }

    // xの親を返す + 経路圧縮
    int find(int x) {
        if (par[x] == x) {
            return x;
        } else {
            return par[x] = find(par[x]);
        }
    }
    
    // xとyが同じ木に属するか判定
    bool same(int x, int y) {
        return find(x) == find(y);
    }

    // xとyを含む木を結合する
    void unite(int x, int y) {
        x = find(x);
        y = find(y);

        if (x == y) return;

        // rank[x] >= rank[y]にする
        if (rank[x] < rank[y]) swap(x, y);
        
        // rankの大きい方、つまりxにyをくっつける
        par[y] = x;
        num[x] += num[y];
        if (rank[x] == rank[y]) rank[x]++;
    }
    
    // xを含む木の頂点数を返す
    int size(int x) {
        return num[find(x)];
    }
};

オブジェクト指向を使った至って無難な実装かと思われます。

では各要素をイジっていきましょう。

(以降説明の都合上、「結合回数」を「時刻」による比喩で表すことが多々あります。察してください。)

メンバ変数

まず、部分永続UFではメンバ変数が2つ増えます。それが主にnowtimeと呼ばれるものです。

nowはint型で、「現在いくつの結合が行われたか」、 すなわち「現在時刻」を保持します。

timeはint型の配列で、各頂点について「直接の親(つまりpar[x])が何度目の結合で更新されたか」を保持します。

通常のUFでは経路圧縮によって直接の親が結合後も変わりますが、経路圧縮を排除することで「一度親が更新されたら永遠にそのまま」という性質ができます。この性質から、時刻time[x]以降par[x]に変更がなされていないことが保障できるのです。

さらにもう1つ、変更が加えられるメンバ変数がnumです。今のnumでは現在の頂点数しか返せませんが、「時刻」と「その時点での頂点数」によるpairをvectorによって保持することで、各時点での要素数を求められるようになります。詳しくは後半の章にて。

というわけで、メンバ変数とコンストラクタが以下のように変わります。

    int now;         // 現在時刻
    int par[V_NUM];  // 各頂点の親
    int rank[V_NUM]; // その頂点を根とする木の深さ
    int time[V_NUM]; // 親がいつ更新されたか
    vector<pair<int, int>> num[V_NUM];
        // (時刻, 頂点数)を要素にもつvector
    
    explicit persistentUF() {
        now = 0;
        for (int i = 0; i < V_NUM; ++i) {
            par[i] = i;
            num[i].push_back(make_pair(0, 1));
                // 時刻0にて頂点数は1
        }
        fill(rank, rank + V_NUM, 0);
        fill(time, time + V_NUM, INF);
            // 自身が根の間は便宜上INFとする
    }

find関数

次にfind関数です。ここが部分永続UFに変更する上での最重要ポイントですが、ことは至ってシンプルです。

find関数は当然ながら頂点xだけでなく、調べたい時刻tを引数に取ります。 そこで、ttime[x]の大小によって場合分けをします。

t < time[x]の場合

この場合、timeの定義から「par[x]は時刻tにてまだ更新されていない」、つまり「時刻tにてxは木の根である」ということが分かります。よってxをそのまま返せばいいです。

t >= time[x]の場合

この場合、「par[x]は時刻tにてすでに更新されている」ため、通常のUF同様その親について再帰的に探索すればOKです。

以上で全パターンを網羅できました。というわけでfind関数と、ついでにsame関数も書き換えましょう。

    // 時刻tにおけるxの親を返す
    int find(int x, int t) {
        if (t < time[x]) {
            return x;
        } else {
            return find(par[x], t);
        }
    }
    
    // 時刻tにてxとyが同じ木に属するか判定
    bool same(int x, int y, int t) {
        return find(x, t) == find(y, t);
    }

unite関数

こちらはUFと大差なく、ただnowtimeの更新が入るだけです。

ただし重要な点が1つ。部分永続UFでは経路圧縮を行わないため、rankによってバランスを保たなければ最悪計算量が O(|V|)になってしまいます。UFをrankなしで実装した方も、今回はちゃんとrankつきで実装しなければなりません。頑張りましょう。

    // 頂点xとyを繋げる
    void unite(int x, int y) {
        ++now; // 時間を進める
        
        x = find(x, now);
        y = find(y, now);

        if (x == y) return;

        if (rank[x] < rank[y]) swap(x, y);
        
        par[y] = x;
        time[y] = now; // 更新時刻を記録
        if (rank[x] == rank[y]) ++rank[x];
    }

ほぼ変わっていないのが分かるでしょう。

しかし、よく見ていただければ分かる通りnumの更新は省いてあります。これは次章にて追加します。

なお「各グループの要素数までは要らないよ」という方は以上で実装終了なので、次の章は丸々飛ばしていただいて構いません。お疲れ様でした。

size関数

正直一番変更が面倒なのがsize関数です。まぁやるしかないんですが。

その前にメンバ変数numについて補足を。例えば2つの木を結合したとき、その2つの木の頂点全てに対して(時刻, 頂点数)のpairを突っ込んでいては計算量的にも空間量的にもアウトです。そのため、numを更新するのは「新しくできた木の根」のみとなります。

これに基づいて、上のunite関数にnumの更新を加えましょう。

    // 頂点xとyを繋げる
    void unite(int x, int y) {
        ++now;
        
        x = find(x, now);
        y = find(y, now);

        if (x == y) return;

        if (rank[x] < rank[y]) swap(x, y);
        
        // parの更新を先にやるとバグるので注意
        num[x].push_back(make_pair(now, size(x, now) + size(y, now)));
        
        par[y] = x;
        time[y] = now;
        if (rank[x] == rank[y]) rank[x]++;
    }

上のような感じになります。コメントにある通り、numの更新位置には気をつけましょう1

さて、本題のsize関数です。まず先程説明した通り、numが更新されるのは木の根に対してのみなので、まずはfind関数で木の根に行く必要があります。

num[x]の中には(時刻, 要素数)のpairが格納されているわけですが、この中には時刻が古すぎるものもあれば、逆に「知りたい時刻」より未来の情報もあります。 その中で適切な情報はどれかというと、「知りたい時刻」と同時刻か、あるいはそれより古い中で最も新しいものです。 num[x]は時刻について昇順に格納されているため、二分探索を使えば高速に適切な情報の入ったindexを調べることができます。

というわけで、少々複雑ですが実装していきましょう。

    // 時刻tにおいて、頂点xを含む木の要素数を返す
    int size(int x, int t) {
        x = find(x, t);
        
        // 適切な情報が入ったindexを二分探索で探り当てる
        int ok = 0, ng = num[x].size();
        while (ng - ok > 1) {
            int mid = (ok + ng) / 2;
            if (num[x][mid].first <= t) {
                ok = mid;
            } else {
                ng = mid;
            }
        }
        
        return num[x][ok].second;
    }

以上で全体の実装が完了です!本当にお疲れ様でした。

全体の実装例

ここまで実装してきたものをまとめると、以下のようになります。

// 頂点数
const int V_NUM = 100010;

class persistentUF {
public:
    int now;         // 現在時刻
    int par[V_NUM];  // 各頂点の親
    int rank[V_NUM]; // その頂点を根とする木の深さ
    int time[V_NUM]; // 親がいつ更新されたか
    vector<pair<int, int>> num[V_NUM];
        // (時刻, 頂点数)を要素にもつvector
    
    explicit persistentUF() {
        now = 0;
        for (int i = 0; i < V_NUM; ++i) {
            par[i] = i;
            num[i].push_back(make_pair(0, 1));
        }
        fill(rank, rank + V_NUM, 0);
        fill(time, time + V_NUM, INF);
    }
    
    // 時刻tにおけるxの親を返す
    int find(int x, int t) {
        if (t < time[x]) {
            return x;
        } else {
            return find(par[x], t);
        }
    }
    
    // 時刻tにてxとyが同じ木に属するか判定
    bool same(int x, int y, int t) {
        return find(x, t) == find(y, t);
    }
    
    // 頂点xとyを繋げる
    void unite(int x, int y) {
        ++now;
      
        x = find(x, now);
        y = find(y, now);

        if (x == y) return;

        if (rank[x] < rank[y]) swap(x, y);
        
        num[x].push_back(make_pair(now, size(x, now) + size(y, now)));
        
        par[y] = x;
        time[y] = now;
        if (rank[x] == rank[y]) ++rank[x];
    }
    
    // 時刻tにおいて、頂点xを含む木の要素数を返す
    int size(int x, int t) {
        x = find(x, t);
        
        int ok = 0, ng = num[x].size();
        while (ng - ok > 1) {
            int mid = (ok + ng) / 2;
            if (num[x][mid].first <= t) {
                ok = mid;
            } else {
                ng = mid;
            }
        }
        
        return num[x][ok].second;
    }
};

とある問題に試しに投げたら無事通ったので、バグはないはずです。

ついでに要素数を省いたものも貼っておきます。

const int INF = 1 << 25;

// 頂点数
const int V_NUM = 100000;

class persistentUF {
public:
    int now;         // 現在時刻
    int par[V_NUM];  // 各頂点の親
    int rank[V_NUM]; // その頂点を根とする木の深さ
    int time[V_NUM]; // 親がいつ更新されたか
    
    explicit persistentUF() {
        now = 0;
        for (int i = 0; i < V_NUM; ++i) {
            par[i] = i;
        }
        fill(rank, rank + V_NUM, 0);
        fill(time, time + V_NUM, INF);
    }
    
    // 時刻tにおけるxの親を返す
    int find(int x, int t) {
        if (t < time[x]) {
            return x;
        } else {
            return find(par[x], t);
        }
    }
    
    // 時刻tにてxとyが同じ木に属するか判定
    bool same(int x, int y, int t) {
        return find(x, t) == find(y, t);
    }
    
    // 頂点xとyを繋げる
    void unite(int x, int y) {
        ++now;
        
        x = find(x, now);
        y = find(y, now);

        if (x == y) return;

        if (rank[x] < rank[y]) swap(x, y);
        
        par[y] = x;
        time[y] = now;
        if (rank[x] == rank[y]) ++rank[x];
    }
};

やはり幾分スッキリしますね。なおこちらはテストしてませんが、不要な部分を切り取っただけなので多分問題はないでしょう。

まとめ

部分永続Union-Findはいいぞ。 今のところ使う機会ほとんどないけどね。


  1. num[x].push_back(make_pair(now, size(x, now - 1) + size(y, now - 1)))のように丁寧に処理すれば更新位置はどこでも問題ありません。