Pythonで2次元バイナリ行列内の全てが1の正方形部分行列の総数を求める方法

2次元のバイナリ行列(0と1のみで構成される行列)が与えられたとき、すべての要素が1である正方形の部分行列が合計いくつ含まれているかを求める問題を考えてみましょう。

入力例

この場合の出力は 17 となります。内訳は、1×1の正方形が12個、2×2の正方形が4個、そして3×3の正方形が1個存在するためです。

解き方:動的計画法(DP)

この問題は、動的計画法を用いることで効率的に解けます。基本的な考え方は、各セルについて「そのセルを右下の角とする最大の正方形の一辺の長さ」を順番に計算していくというものです。

重要なポイントとして、あるセルに記録された値は「その位置を右下とする正方形の個数」そのものと一致します。したがって、すべてのセルの値を合計すれば、全体の正方形の総数が求まります。

具体的な手順は以下の通りです。

• 結果を格納する変数 res を0で初期化する
• i を0から行数までループする
  • j を0から列数までループする
    • i が0、または j が0の場合(最上行または最左列)
      • res := res + matrix[i][j](そのまま加算)
    • それ以外で matrix[i][j] が1の場合
      • matrix[i][j] = min(matrix[i][j-1], matrix[i-1][j], matrix[i-1][j-1]) + 1(左・上・左上の隣接セルの最小値に1を加える)
      • res := res + matrix[i][j]
• 最後に res を返す

なぜこの式が成り立つのかというと、あるセルが1であるとき、そこを右下とする正方形の大きさは、その左隣・上隣・左上のセルそれぞれが持つ正方形のサイズに制約されるためです。3つの隣接セルの最小値に1を足すことで、現在のセルにおける最大の正方形サイズが確定します。

実装例

class Solution:
   def solve(self, matrix):
      res = 0
      for i in range(len(matrix)):
         for j in range(len(matrix[0])):
            if i == 0 or j == 0:
               res += matrix[i][j]
            elif matrix[i][j] == 1:
               matrix[i][j] = min(matrix[i][j - 1], matrix[i - 1][j], matrix[i - 1][j - 1]) + 1
               res += matrix[i][j]
      return res

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

入力

matrix = [  
[1, 1, 1, 0],  
[1, 1, 1, 0],  
[1, 1, 1, 0],  
[0, 0, 0, 0],  
[1, 0, 1, 1]
]

出力

17

まとめ

この手法の計算量は O(行数 × 列数) であり、行列内の全セルを一度走査するだけで済むため非常に効率的です。全てが1の正方形部分行列を数える問題は、LeetCodeなどの競技プログラミングでも頻出のテーマなので、ぜひ覚えておきましょう。