この記事について
Union-findについて勉強したので、LeetCodeのNumber of Islandsという問題をUnion-findを使用して解いてみます。
Union-findとは？
Union-findは素集合データ構造を操作するためのアルゴリズムで、2つの要素が同じ集合に属しているのかを判定したり、2つの集合をまとめたりする操作を行います。
私もUnion-findは知らなかったのですが、AtCoderが提供しているスライドが非常にわかりやすかったです。
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です。陸の集合の個数は、陸の集合の根の個数(上述のスライド参照)と等しいので、根の個数を数えてあげれば終了です。
ざっくりと疑似コードを示すと次のようになります。
- 地図の各グリッドに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程度で走るので、まあまあ悪くない数字なんじゃないでしょうか。
