Pythonで1のみからなる正方形の部分行列の数を数える方法（動的計画法）

問題の概要

m × n の2値（バイナリ）行列が与えられたとき、すべての要素が 1 である正方形の部分行列がいくつ存在するかを求めます。

例として、次のような入力を考えてみましょう。

この場合、出力は 15 になります。内訳は以下のとおりです。

• 1辺が1の正方形：10個
• 1辺が2の正方形：4個
• 1辺が3の正方形：1個

合計：10 + 4 + 1 = 15

解法のアプローチ（動的計画法）

この問題は、動的計画法（DP）を使うことで効率的に解けます。各セルについて「そのセルを右下とする最大の正方形の一辺の長さ」を記録していくのがポイントです。

手順は以下のとおりです。

1. 行列が [[1]] のみの場合は、答えは 1 を返します。
2. 行数 rows と列数 cols を取得し、結果格納用の変数 result を 0 で初期化します。
3. すべてのセルを走査します。
   • 最上行または最左列（row == 0 または col == 0）の場合、そのセルが 1 なら result に 1 を加算します（それより大きい正方形は作れないため）。
   • それ以外のセルで値が 1 の場合、左・上・左上の3つのセルの最小値に 1 を加えた値を square とし、その値を現在のセルに書き込みます。そして result に square を加算します。
4. 最後に result を返します。

この方法では、各セルを一度だけ訪れるため、時間計算量は O(m×n) となり、全ての正方形を総当たりで調べる方法よりも大幅に効率化できます。

Pythonでの実装例

def solve(matrix):
    if matrix == [[1]]:
        return 1
    rows = len(matrix)
    cols = len(matrix[0])
    result = 0
    for row in range(rows):
        for col in range(cols):
            if (row == 0 or col == 0):
                if matrix[row][col] == 1:
                    result += 1
            elif matrix[row][col] == 1:
                square = min(matrix[row-1][col], min(matrix[row][col-1], matrix[row-1][col-1])) + 1
                matrix[row][col] = square
                result += square
    return result

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

入力

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

出力

15

コードのポイント

• DPテーブルの再利用： 元の行列 matrix 自体をDPテーブルとして上書きすることで、追加のメモリを削減しています。
• 累積カウント： 各セルの square の値は「そのセルを右下とする正方形の個数」を表すため、単純に加算するだけで全体の個数が求まります。
• 境界の扱い： 最上行・最左列は正方形を拡張できないため、1であれば常に1だけカウントします。

このアルゴリズムを使えば、大きな行列でも線形時間で正方形部分行列の総数を正確に数えることができます。