けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder Library Practice Contest A - Disjoint Set Union

Union-Find の最も典型的な練習問題

問題概要

 N 頂点  0 辺の無向グラフに  Q 個のクエリが飛んでくる

  •  0, u, v:辺  (u,v) を追加する
  •  1, u, v 2 頂点間  u,v が連結ならば 1、そうでないなら 0 を出力する

制約

  •  1 \le N, Q \le 2 \times 10^{5}

解法

Union-Find はグループ分けを管理するデータ構造です。以下のクエリをサポートします。

  • merge(u, v) u を含むグループと、 v を含むグループを統合して 1 つの大きなグループを作る
  • same(u, v) u, v が同一のグループに属するかどうかを判定する

 N 個の要素がそれぞれ単独で存在している状態から開始して、クエリ merge(u, v) が呼ばれるたびにグループが大きくなっていきます。

たとえば上図左のように {0, 2, 4, 7}, {3, 5}, {6} という 3 つのグループがある状態で、merge(2, 3) を実行すると、上図右のように {0, 2, 3, 4, 5, 7}, {6} という 2 つのグループになります。

この Union-Find を用いることで、各クエリに要する計算量は  O(\alpha(N)) であり、全体の計算量は  O(N + Q\alpha(N)) となります。

ここで、 \alpha(N) はアッカーマン関数の逆関数と呼ばれるもので (知らなくても問題ないです)、とても小さな値になることが知られています。 N \le 10^{80} に対しては  \alpha(N) \le 4 となります!!

コード

Union-Find は ACL において、atcoder::dsu として実装されています。

#include <bits/stdc++.h>
#include <atcoder/dsu>
using namespace std;

int main() {
    int N, Q;
    cin >> N >> Q;

    // サイズ N の Union-Find を宣言
    atcoder::dsu uf(N);

    // クエリ処理
    for (int i = 0; i < Q; ++i) {
        int type, u, v;
        cin >> type >> u >> v;
        if (type == 1) 
            cout << (uf.same(u, v) ? 1 : 0) << endl;
        else
            uf.merge(u, v);
    }
}

自前ライブラリでも AC

#include <bits/stdc++.h>
using namespace std;

struct UnionFind {
    vector<int> par;

    UnionFind() { }
    UnionFind(int n) : par(n, -1) { }
    void init(int n) { par.assign(n, -1); }
    
    int root(int x) {
        if (par[x] < 0) return x;
        else return par[x] = root(par[x]);
    }
    
    bool issame(int x, int y) {
        return root(x) == root(y);
    }
    
    bool merge(int x, int y) {
        x = root(x); y = root(y);
        if (x == y) return false;
        if (par[x] > par[y]) swap(x, y); // merge technique
        par[x] += par[y];
        par[y] = x;
        return true;
    }
    
    int size(int x) {
        return -par[root(x)];
    }
};

int main() {
    int N, Q;
    cin >> N >> Q;

    // サイズ N の Union-Find を宣言
    UnionFind uf(N);

    // クエリ処理
    for (int i = 0; i < Q; ++i) {
        int type, u, v;
        cin >> type >> u >> v;
        if (type == 1) 
            cout << (uf.issame(u, v) ? 1 : 0) << endl;
        else
            uf.merge(u, v);
    }
}