LeetCode 200: Number of Islands をUnion-findを使って解く

この記事について

Union-findについて勉強したので、LeetCodeのNumber of Islandsという問題をUnion-findを使用して解いてみます。

Union-findとは？

Union-findは素集合データ構造を操作するためのアルゴリズムで、2つの要素が同じ集合に属しているのかを判定したり、2つの集合をまとめたりする操作を行います。

私もUnion-findは知らなかったのですが、AtCoderが提供しているスライドが非常にわかりやすかったです。

Union find(素集合データ構造) from AtCoder Inc.

www.slideshare.net

LeetCode 200: Number of Islands

Union-findを理解したところで、早速問題を見てみます。問題はLeetCodeのものをお借りしています。
LeetCode: コーディング面接に向けた練習に使えるサイトの紹介 - 今日も窓辺でプログラム

'1'(陸)と'0'(海)で構成されている2次元のグリッド地図が与えられるので、島の個数を求めよ。島は海に囲まれていて、縦横につながっている陸で作られるものとする。地図の四辺の外側は海であると仮定してよい。
https://leetcode.com/problems/number-of-islands/

また問題文には入力の例として次の2つが示されています。

例1：

島の個数 = 1

例2:

島の個数 = 3

この問題、Union-findの考え方を使って解くことができますので、よかったら挑戦してみてください。

Union-findを使った回答

例2を見てみます。

この入力の島の数は3です。左上の4つの'1'の島、中央の1つの'1'の島、そして右下の2つの'1'の島です。
つまり、上下左右でつながっている陸('1')を順にまとめていって、最終的に残った陸の集合の数が島の数となります。
この、陸をまとめる部分で、Union-findのunionの考え方が適用できます。

全ての陸がまとまったあとは、陸の集合がいくつあるかを数えればOKです。陸の集合の個数は、陸の集合の根の個数(上述のスライド参照)と等しいので、根の個数を数えてあげれば終了です。

ざっくりと疑似コードを示すと次のようになります。

地図の各グリッドにIDを割り振る。海(='0')に関しては全て-1を、陸(='1')に関しては0以上の整数を一意に与える。
グリッドを横向きと縦向きに走査し、陸地が隣り合っている場合は2つの陸地を1つにまとめる (union)
陸の集合の個数を、その根の個数を数えることによって求める

私の実装例はこんな感じになります。
ステップ1でuidという配列にIDを割り当てていき、uidに対してunionやfindの操作を行っています。

class Solution {
public:
    // Union-findのfindの定義。uid[n]の根を求める。
    int findRoot(vector<int>& uid, int n) {
        if (uid[n] == n) {
            return n;
        } else {
            // 経路圧縮の処理。ランクを使った効率化は今回は実装しない。
            uid[n] = findRoot(uid, uid[n]);
            return uid[n];
        }
    }

    // Union-findのUnionの定義。xとyを一つの集合にまとめる。
    void unify(vector<int>& uid, int x, int y) {
        x = findRoot(uid, x);
        y = findRoot(uid, y);

        if (x == y) { return; }
        uid[y] = x;
    }

    int numIslands(vector<vector<char>>& grid) {
        if ((grid.size() == 0) || (grid[0].size() == 0)) { return 0; }

        auto m = grid.size();
        auto n = grid[0].size();
        vector<int> uid(m * n, -1); 

        // ステップ1：uidに初期IDを割り当てる。
        // 海は全て-1で、陸にはその位置のインデックスをもとに0以上のIDを与える。
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                if (grid[i][j] == '1') {
                    uid[i * n + j] = i * n + j;
                }
            }
        }

        // ステップ2-1: 横方向に陸が2つ続いている場合は、それらを一つの集合にまとめる
        for (int i = 0; i < m; i++) {
            for (int j = 1; j < n; j++) {
                if ((grid[i][j - 1] == '1') && (grid[i][j] == '1')) {
                    unify(uid, i * n + j - 1, i * n + j);
                }
            }
        }

        // ステップ2-2: 縦方向に陸地が2つ続いている場合は、それらを一つの集合にまとめる
        for (int j = 0; j < n; j++) {
            for (int i = 1; i < m; i++) {
                if ((grid[i - 1][j] == '1') && (grid[i][j] == '1')) {
                    unify(uid, (i - 1) * n + j, i * n + j);
                }
            }
        }

        // ステップ3: 陸の集合の根の個数を数える。
        // IDが0以上だと陸であり、なおかつ陸のインデックスとIDが等しいものが陸の集合の根である。
        int ans = 0;
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                auto idx = i * n + j;
                if ((uid[idx] >= 0) && (uid[idx] == idx)) {
                    ans++;
                }
            }
        }

        return ans;
    }
};

先ほどのスライドによると、経路圧縮のみを実装した場合はfindやunionの計算量はO(logN)となるようですので、それを信じるとこのアルゴリズムの時間計算量はO(N logN)となります。
このコードを提出すると6ms程度で走るので、まあまあ悪くない数字なんじゃないでしょうか。