Pythonで全ての要素が1の部分行列を数えるプログラム(DP活用)

m × n のバイナリ行列(各要素が 0 または 1 の行列)が与えられたとき、「すべての要素が 1」で構成される部分行列がいくつあるかを求める問題を考えます。

例として、次のような 3 × 3 の行列を入力してみましょう。

この場合の出力は 13 になります。内訳は次の通りです。

• 1 × 1 の部分行列:6 個
• 2 × 1 の部分行列:3 個
• 1 × 2 の部分行列:2 個
• 3 × 1 の部分行列:1 個
• 2 × 2 の部分行列:1 個

解き方のアプローチ

この問題は、動的計画法(DP)の考え方を使うことで効率よく解けます。手順は以下の通りです。

• m := 行列の行数とする
• n := 行列の列数とする
• dp := 同じサイズ(m × n)のゼロ行列を用意する
• i を 0 から m − 1 まで繰り返す:
    j を 0 から n − 1 まで繰り返す:
      i = 0 かつ matrix[i][j] が 1 の場合 → dp[i][j] := 1
      それ以外で matrix[i][j] が 1 の場合 → dp[i][j] := dp[i−1][j] + 1
• total := 0 とする
• i を 0 から m − 1 まで繰り返す:
    j を 0 から n − 1 まで繰り返す:
      k を j + 1 から n まで繰り返す:
        total := total + min(dp[i][j:k])
• total を返す

なぜこれで正しく数えられるのか

dp[i][j] には「セル (i, j) を下端としたとき、その列を上方向に何マス連続で 1 が続いているか」という高さが記録されます。ある行 i を底辺とし、列の範囲 [j, k) を横幅として選んだとき、そこに作れる全て 1 の部分行列の個数は、範囲内の高さの最小値 min(dp[i][j:k]) に等しくなります。すべての行・すべての列範囲についてこの値を合計すれば、条件を満たす部分行列の総数が求まります。

Pythonでの実装例

def solve(matrix):
    m = len(matrix)
    n = len(matrix[0])
    dp = [[0] * n for _ in range(m)]
    for i in range(m):
        for j in range(n):
            if i == 0 and matrix[i][j]:
                dp[i][j] = 1
            elif matrix[i][j]:
                dp[i][j] = dp[i-1][j] + 1
    total = 0
    for i in range(m):
        for j in range(n):
            for k in range(j+1, n+1):
                total += min(dp[i][j:k])
    return total

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

入力

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

出力

13

計算量について

DPテーブルの構築には O(mn) かかります。集計処理では各行について列範囲のすべての組み合わせに対して最小値を求めるため、全体の計算量は O(m · n²) となります。より大きな入力に対応したい場合は、各行の高さ配列に対して単調スタック(monotonic stack)を用いることで、O(mn) まで高速化できます。