Union Find Treeを使う(互いに素な集合)

Union Find Tree: 互いに素な集合

この記事について

Aizu Online Judgeに掲載されている問題をC++(大体C++11)で解くという遊びをしているので、考え方の例と解けたコード例を記録していこうかと思います。コード例はたたんでおきます。クリックやタップでコード例を開けます。

こちらの問題:

Aizu Online Judge

1つのクエリあたり $O(n)$ で解決する方法で最初はやってしまいましたが、調べると適したデータ構造があるようなのでそれで実装し直しました。

2023/10追記: 類題:

012 - Red Painting(★4)
AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.

DSL_1_A: Disjoint Set: 互いに素な集合

同じ要素を複数の集合が持たないように、要素による集合の併合と探索を行う問題。英語での問題集のタイトルから何かデータ構造があるんだなとは思いましたが、最初はなるべく自力でやるかということでなんとかしてみました。

考えたこと(効率悪い)

これは効率があまり良くない方法です。 $n = 10^8$ などになってくると時間がかかり過ぎてしまうと思います。

$[0,n-1]$の整数を使うので、集合を区別するために、サイズ $n$ の配列を使ってインデックスを各要素として扱い、配列の各値を所属している集合の代表のインデックスにしておく、という方法を考えました。例えば、A[7]A[1]と同じ集合に属すときはA[7]==1となります。

代表のインデックスというのは、今回は、各集合で最も小さいインデックスを格納するようにしました。

例としては、 n=6で、 $0\ 1\ 2\ 3\ 4\ 5$ が初期状態とすると、クエリunite 1 3の実行後は、 $0\ 1\ 2\ 1\ 4\ 5$ となります。

これだと、1つのuniteに対して、配列全てにアクセスして片方の同じ集合だった要素の値をもう片方の集合の代表インデックスに変更する必要があります。

$O(n+q)$ くらい?(qはクエリの数)だし、まあいいかと思って提出してみると、パスしてしまいました。しかし、若干時間がかかっているようなので実装し直しが必要かなと思ったので書き直しました。

効率悪めな実装例

一応折角なので効率の悪いもののコード例を載せておきたいと思います。 lvrvで確保しておかないと、ループ中にA[l]A[r]が置き換わってしまうことを忘れてて若干詰まってました。

折り畳んであります。
// Disjoint Set

#include <iostream>
#include <vector>

int main()
{
    int n, q;
    std::cin >> n >> q;
    std::vector<int> A(n);
    for (size_t i = 0; i < n; i++)
    {
        A[i] = i;
    }
    for (size_t j = 0; j < q; j++)
    {
        int comp, l, r;
        std::cin >> comp >> l >> r;
        if (comp == 0)
        {
            // 小さい方に合わせるため、それを取得
            int m = std::min(A[l], A[r]);
            // ループ中に書き換えられるので、予め取得しておく
            int lv = A[l];
            int rv = A[r];
            for (size_t i = 0; i < n; i++)
            {
                // 該当要素(両集合の要素)の値を最小値で揃える
                if ((A[i] == lv) || (A[i] == rv))
                    A[i] = m;
            }
        }
        else
        {
            // sameは2つのインデックスの値をチェックするだけ
            std::cout << (A[l] == A[r] ? 1 : 0) << std::endl;
        }
    }
    return 0;
}

効率の良い方法

重複要素のない集合の併合・探索に適したUnion Find Treeというデータ構造があるそうです。

とりあえず見つけた資料(PDF)はこちら: Union-Find Algorithms

配列を用意して、同一集合のインデックスの値を持たせるのは同じようですが、木構造を意識することにより、unite操作を効率よく行えるようです。

配列の値には、木構造として捉えたときの親のインデックスを使います。再帰的に辿って行けば根を見つけられます。根が同じ値なら、同じ集合に属していることになります。これでsame操作を行えます。

本実装では、unite操作を行うときはxy双方の根を見つけて、大きいインデックスの方を小さいインデックスの方の子になるように繋ぎ直します。別にどういう方針で繋ぎ直してもいいです。本当は、小さい木を大きい木に繋ぎ直したほうが効率は良くなりますが、ここまでしなくても良いパフォーマンスが出たので終いとしました。

しかし、ただuniteで木を繋ぎ直していくと、木の高さが大きくなる可能性があります。最大で $n$ になります。これでは根を見つける操作で $O(n)$ になるのでよくありません。しかもこの操作はsameuniteどちらでも使うので場合によっては最初の実装よりも悪化しそうです。

木の高さを抑えるには部分木を根の直下に持ってくればいいので、根を見つける操作の際、ついでに辿ったノードの親を見つけた根に繋ぎ直すことで簡単に実現できます。

効率改善した実装例

繋ぎ直しのサイズ考慮はせずに高さを圧縮する(経路圧縮する)場合の実装になります。

折り畳んであります。
// union find tree

#include <iostream>
#include <vector>

// x,yが所属する集合が同じかどうかチェックする
bool same(std::vector<int> *A, int x, int y);
// x,yが所属する木を併合する
void unite(std::vector<int> *A, int x, int y);
// iが所属する木の根を見つける
int root(std::vector<int> *A, int i);

int main()
{
    int n, q;
    std::cin >> n >> q;
    std::vector<int> A(n);
    for (size_t i = 0; i < n; i++)
    {
        A[i] = i;
    }

    for (size_t j = 0; j < q; j++)
    {
        int op, x, y;
        std::cin >> op >> x >> y;
        if (op == 0)
        {
            unite(&A, x, y);
        }
        else
        {
            std::cout << (same(&A, x, y) ? 1 : 0) << std::endl;
        }
    }

    return 0;
}

bool same(std::vector<int> *A, int x, int y)
{
    // 根が同じなら所属する集合は同じ
    return root(A, x) == root(A, y);
}

void unite(std::vector<int> *A, int x, int y)
{
    int rootx = root(A, x);
    int rooty = root(A, y);
    if (rootx == rooty)
        return;
    // インデックスが小さい方を根にするように繋ぎかえる
    if (rootx < rooty)
        (*A)[rooty] = rootx;
    else
        (*A)[rootx] = rooty;
}

int root(std::vector<int> *A, int i)
{
    int parent = (*A)[i];
    if (i == parent)
        return i;
    // 根でないときは再帰する。戻ったら根を親にする
    (*A)[i] = root(A, parent);
    return (*A)[i];
}

おわり

小さい木と大きい木を判別するにはサイズを保持する配列が新しく必要になりそうです。と思いましたが、根が持つ値をサイズ+nなどにしておけば、根の判別とサイズの取得の両方をメモリを節約しつつ達成できそうです。と思ったけど、木の高さが2になることがあるので、やっぱりサイズの取得にちょっとコストがかかるかな。

効率の良くないやり方をpythonなどのインタプリタ言語でやっても1クエリあたり $O(n)$ だからパスはするかと思います。

計算量の求め方がいまいち分かってないのがちょっと心残りです…

木構造って便利。以上です。


Amazonアソシエイト

参考にしたPDFの資料と著者が同じかも?どちらもプリンストン大学の方のようです。

Bitly

ネットでも読めるようです。

Case Study: Union-Find
The textbook Algorithms, 4th Edition by Robert Sedgewick and Kevin Wayne surveys the most important algorithms and data ...

コメント

タイトルとURLをコピーしました