Pythonで1つの水のセルを陸地に変更した後、最大の島を見つけるプログラム

問題の概要

1が陸地、0が水を表す2次元のバイナリ行列(マトリックス)が与えられます。「島」とは、上下左右につながった1の集まりであり、周囲を水に囲まれた領域のことです。この問題では、最大で1つの水のセルを陸地のセルに変更できるという条件のもとで、作り出せる最大の島のサイズを求める必要があります。

入力例

この場合の出力は 7 になります。たとえば左端の2行目にある水のセルを陸地に変更すると、左上の小さな島(サイズ1)と下側の大きな島(サイズ5)がつながり、合計7の島が完成するからです。変更後の行列は次のようになります。

アルゴリズムの考え方

すべての水のセルを順番に陸地へ変えて、そのたびに島のサイズを数え直す方法もありますが、これは非常に非効率です。そこで、フラッドフィル(洪水填充)を使って各島に一意なIDを振り、あらかじめ面積を記録しておくことで効率よく解きます。手順は以下の通りです。

• R := 行列の行数、C := 行列の列数とします。
• mass := 島のIDとその面積を記録する新しい辞書(マップ)を用意します。
• id := 555 とします。0や1と衝突しない一意なIDとして使うためです。
• floodfill(r, c, id) 関数を定義します。
  • r と c が行列の範囲内で、mat[r][c] が 1 の場合:
    • mat[r][c] を id に書き換えます。
    • mass[id] を 1 増やします。
    • 上下左右の4方向 (r+1,c)、(r−1,c)、(r,c+1)、(r,c−1) に対して再帰的に floodfill を呼び出します。
• メイン処理では、まず全マスを走査し、値が 1 のマスを見つけたら id を1増やして mass[id] = 0 と初期化し、floodfill でその島全体に同じIDを振って面積を記録します。
• ans := mass のすべての値と 1 の最大値とします(島が1つも存在しないケースでも最低 1 を返すため)。
• 再度全マスを走査し、水のセル(0)ごとに次を行います。
  • 上下左右の隣接セルを調べ、範囲内かつ陸地であれば、その島のIDを island_set(集合)に追加します。集合を使うことで、同じ島が複数方向に隣接していても重複カウントされません。
  • ans := max(ans, 1 + island_set 内の各島の面積の合計) で更新します。
• 最後に ans を返します。

この方法なら、時間計算量 O(R×C)、空間計算量 O(R×C) で問題を解くことができます。

実装例(Python)

class Solution:
    def solve(self, mat):
        R, C = len(mat), len(mat[0])
        mass = {}
        id = 555

        def floodfill(r, c, id):
            nonlocal R, C, mat, mass
            if 0 <= r < R and 0 <= c < C and mat[r][c] == 1:
                mat[r][c] = id
                mass[id] += 1
                for r2, c2 in [(r + 1, c), (r - 1, c),
                                (r, c + 1), (r, c - 1)]:
                    floodfill(r2, c2, id)

        for r in range(R):
            for c in range(C):
                if mat[r][c] == 1:
                    id += 1
                    mass[id] = 0
                    floodfill(r, c, id)

        ans = max(list(mass.values()) + [1])
        for r in range(R):
            for c in range(C):
                if mat[r][c] != 0:
                    continue
                island_set = set()
                for r2, c2 in [(r + 1, c), (r - 1, c),
                                (r, c + 1), (r, c - 1)]:
                    if 0 <= r2 < R and 0 <= c2 < C and mat[r2][c2]:
                        island_set.add(mat[r2][c2])
                ans = max(ans, 1 + sum(mass[island] for island in island_set))
        return ans


ob = Solution()
matrix = [
    [1, 0, 1],
    [0, 0, 0],
    [1, 1, 0],
    [1, 1, 1]
]
print(ob.solve(matrix))

入力

[
[1, 0, 1],
[0, 0, 0],
[1, 1, 0],
[1, 1, 1]
]

出力

7

まとめ

このアルゴリズムのポイントは、フラッドフィルによって各島の面積を事前に記録しておき、水のセルを仮想的に陸地へ変えた場合の結合結果を、実際に行列を書き換えることなく計算できる点にあります。これにより、総当たり的なアプローチよりも大幅に高速に最大の島のサイズを求められます。