Pythonで与えられたバイナリ行列の正方形部分行列の数を数えるプログラム

問題概要

2次元のバイナリ行列(0と1だけで構成された行列)が与えられたとき、すべての要素が1で構成されている正方形部分行列の総数を求めることを考えます。

例えば、次のような入力が与えられたとします。

この場合の出力は 5 になります。これは、(2 × 2) の正方形が1つと、(1 × 1) の正方形が4つ存在するためです。

解法のアプローチ(動的計画法)

この問題は動的計画法(DP)を使うことで効率的に解けます。基本的な考え方は、「各セルを右下の角とする最大の正方形の一辺の長さ」を行列に順番に記録していくというものです。あるセルが1であれば、その左・上・左上のセルの値の最小値に1を足したものが、そのセルから作れる最大の正方形のサイズになります。この値をカウンターに加算していくことで、すべての大きさの正方形を重複なく数えられます。

アルゴリズムの手順

• 行列 mat が空であれば、0 を返します。
• カウンター変数 c を 0 で初期化します。
• i を 0 から行数まで、j を 0 から列数までループします。
  • mat[i][j] が 1 の場合:
    • i が 0 または j が 0 の場合(端の行・列):c に 1 を加算します。
    • それ以外の場合:temp = min(mat[i-1][j-1], mat[i][j-1], mat[i-1][j]) + mat[i][j] を計算し、c に temp を加算したうえで、mat[i][j] を temp で更新します。
• 最後に c を返します。

Python実装例

それでは、実際のコードを見てみましょう。

def solve(mat):
   if mat == []:
      return 0
   c = 0

   for i in range(len(mat)):
      for j in range(len(mat[0])):
         if mat[i][j] == 1:
            if i == 0 or j == 0:
               c += 1
            else:
               temp = min(mat[i - 1][j - 1], mat[i][j - 1], mat[i - 1][j]) + mat[i][j]
               c += temp
               mat[i][j] = temp
   return c

matrix = [
   [0, 1, 1],
   [0, 1, 1]
]
print(solve(matrix))

入力

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

出力

5

計算量について

このアルゴリズムの時間計算量は O(m × n)(m は行数、n は列数)です。また、元の行列をそのまま書き換えて利用しているため、追加のメモリはほとんど不要です。すべての正方形候補を総当たりで確認する非効率な手法と比べると、大幅に高速に動作します。