C++で1だけで構成される部分行列の個数を数えるプログラム

0と1のみから構成される行列が与えられたとき、その中に存在する「すべての要素が1である部分行列」の個数を求める問題を考えてみましょう。求めた個数を出力として表示します。

入力例

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

この場合、すべての要素が1である部分行列は全部で12個存在するため、出力は 12 となります。

解法のアプローチ

この問題を効率的に解くには、二次元累積和(2Dプレフィックスサム) と 二分探索 を組み合わせます。手順は以下の通りです。

• n を行列の行数、m を行列の列数とします。
• (n+1) × (m+1) のサイズを持つ配列 add(二次元累積和)を定義します。
• 各マスに対して累積和を次のように計算します。
  • add[i+1][j+1] += matrix[i][j]
  • add[i+1][j+1] += add[i][j+1]
  • add[i+1][j+1] += add[i+1][j]
  • add[i+1][j+1] -= add[i][j](重複分を引く)
• 答えを格納する変数 res を 0 で初期化します。
• 各行 i、各列 j について以下を繰り返します。
  • matrix[i][j] が 0 であれば、このセルを左上とする部分行列は存在しないため、次の反復へスキップします。
  • 部分行列の高さ k を 1 から (n − i) まで変化させながら、二分探索で「すべてが1となる最大の幅」を求めます。
    • p := 0、q := m − j として初期化します。
    • p ≤ q の間、x := (p + q) / 2 とし、a := k * x(領域の面積)を計算します。
    • cur := add[i+k][j+x] − add[i][j+x] − add[i+k][j] + add[i][j] により、対象領域の合計値を O(1) で取得します。
    • cur == a であれば、その領域はすべて1なので、r := x として探索範囲を右側に広げます(p := x + 1)。そうでなければ q := x − 1 として範囲を狭めます。
  • r == 0 の場合は、これ以上幅を広げられないため内側のループを抜けます。
  • res := res + r として個数を加算していきます。
• 最後に res を返します。

計算量

累積和により任意の矩形領域の合計が定数時間で求められるため、各セル・各高さについて二分探索 O(log m) を行う形になり、素朴な全探索よりも大幅に高速化できます。

C++での実装例

理解を深めるために、実際の実装を見てみましょう。

#include<bits/stdc++.h>
using namespace std;

int solve(vector<vector<int>>& matrix) {
    int n = matrix.size();
    int m = matrix[0].size();
    int add[n + 1][m + 1];
    memset(add, 0, sizeof(add));

    for (int i = 0; i < n; i++) {
        for (int j = 0; j < m; j++) {
            add[i + 1][j + 1] += matrix[i][j];
            add[i + 1][j + 1] += add[i][j + 1];
            add[i + 1][j + 1] += add[i + 1][j];
            add[i + 1][j + 1] -= add[i][j];
        }
    }
    int res = 0;
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < m; j++) {
            if (!matrix[i][j])
                continue;
            for (int k = 1; k <= (n - i); k++) {
                int p = 0,
                    q = m - j;
                int r;
                while (p <= q) {
                    int x = (p + q) / 2;
                    int a = k * x;
                    int cur = add[i + k][j + x] - add[i][j + x] - add[i + k][j] + add[i][j];
                    if (cur == a) {
                        r = x;
                        p = x + 1;
                    } else
                        q = x - 1;
                }
                if (r == 0)
                    break;
                res += r;
            }
        }
    }
    return res;
}
int main() {
    vector<vector<int>> mat = {{0, 0, 1, 0}, {0, 1, 0, 0}, {0, 1, 0, 1}, {1, 1, 0, 1}};
    cout<< solve(mat) <<endl;
    return 0;
}

入力

{{0, 0, 1, 0}, {0, 1, 0, 0}, {0, 1, 0, 1}, {1, 1, 0, 1}}

出力

12