C++で1のみを含む正方形の部分行列を数える方法

m × n のサイズを持つバイナリ行列(0と1のみで構成される行列)が与えられたとき、すべての要素が1である正方形の部分行列の総数を数える問題を考えてみましょう。

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

この場合、答えは 15 となります。内訳は以下の通りです。

• 1 × 1 の正方形(1つの1からなる):10個
• 2 × 2 の正方形(4つの1からなる):4個
• 3 × 3 の正方形(9つの1からなる):1個

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

この問題は、動的計画法を用いることで効率的に解くことができます。基本的な考え方は、「各セルを右下の角とする正方形の個数」をそのセルの値として記録していくというものです。

あるセル matrix[i][j] が1であるとき、そこで終わる正方形の数は、右・下・右下の隣接セルの値の最小値に1を加えたものになります。これは、より大きな正方形を作るためには、周囲の3方向すべてが同じ大きさの正方形を形成している必要があるためです。

アルゴリズムの手順

1. 答えを格納する変数 ans を0で初期化し、行数を n、列数を m とします。
2. 最下行の各セルの値を ans に加算します(最下行のセルは1×1の正方形としてしか使えないため)。
3. 最右列の各セルの値を ans に加算します。
4. この際、右下隅のセルが二重にカウントされているため、matrix[n-1][m-1] を一度差し引きます。
5. 残りのセルについて、i を n-2 から 0 へ、j を m-2 から 0 へと逆順に走査します。
   • matrix[i][j] == 1 の場合:matrix[i][j] = 1 + min(matrix[i+1][j+1], matrix[i][j+1], matrix[i+1][j]) と更新します。
   • それ以外の場合:matrix[i][j] = 0 とします。
   • 更新後の値を ans に加算します。
6. 最後に ans を返します。

C++による実装例

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
   int countSquares(vector<vector<int>>& matrix) {
      int ans = 0;
      int n = matrix.size();
      int m = matrix[0].size();
      for(int i = 0; i < m; i++)ans += matrix[n-1][i];
      for(int i = 0; i < n; i++)ans += matrix[i][m-1];
      ans -= matrix[n-1][m-1];
      for(int i = n - 2;i >= 0; i--){
         for(int j = m-2 ;j >= 0; j--){
            matrix[i][j] = matrix[i][j] == 1? 1 + min({matrix[i+1][j+1],matrix[i][j+1],matrix[i+1][j]}) : 0;
            ans += matrix[i][j];
         }
      }
      return ans;
   }
};
main(){
   vector<vector<int>> v = {{0,1,1,1},{1,1,1,1},{0,1,1,1}};
   Solution ob;
   cout << (ob.countSquares(v));
}

入力

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

出力

15

計算量の分析

このアルゴリズムの時間計算量は O(m × n) です。行列の各セルを一度だけ訪問するためです。また、入力行列自体をDPテーブルとして再利用しているため、追加の空間計算量は O(1) で済みます(入力を書き換えてもよい場合)。全探索的にすべての正方形候補を調べる方法では O((mn)²) 程度かかる可能性がありますが、このDP手法により大幅な高速化が実現できます。