読者です 読者をやめる 読者になる 読者になる

今日も窓辺でプログラム

外資系企業勤めのエンジニアが勉強した内容をまとめておくブログ

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

この記事について

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

Union-findとは?

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

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

www.slideshare.net

LeetCode 200: Number of Islands

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

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

https://leetcode.com/problems/number-of-islands/

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

例1:

11110
11010
11000
00000

島の個数 = 1

例2:

11000
11000
00100
00011

島の個数 = 3

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

Union-findを使った回答

例2を見てみます。

11000
11000
00100
00011

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

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

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

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

私の実装例はこんな感じになります。
ステップ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程度で走るので、まあまあ悪くない数字なんじゃないでしょうか。