C++で要素の合計が指定した値と等しくなる部分行列の個数を求めるプログラム

整数要素を含む行列が与えられたとき、その中から要素の合計が指定した目標値(target)と等しくなる部分行列をすべて見つけ、その個数を返すことを考えます。

問題の例

例えば、次のような行列が入力だとします。

ここで target = 5 とすると、出力は 3 になります。これは、要素の合計が 5 に等しい部分行列がちょうど 3 つ存在するためです。具体的には、「下2行全体」「上から2〜4行目 × 左から2〜4列目」「すべての行 × 左から3列目まで」の3つが該当します。

解法のアプローチ

この問題は、「1次元配列内で合計が k になる連続部分配列の個数を数える」古典的なテクニック(累積和 + ハッシュマップ)を2次元に拡張することで効率よく解けます。

アルゴリズムの手順

1. 行数 n と列数 m を求めます(n = 0 の場合は m = 0 とします)。
2. m > n の場合は行列を転置してから同じ処理を再帰的に適用します。これにより常に列数 ≤ 行数となり、計算量を最小化できます。
3. 答えを格納する変数 ans を 0 で初期化します。
4. 左端の列 p を 0 から m-1 まで順に固定し、さらに右端の列 q を p から m-1 まで伸ばしながら以下を実行します。
   • 各行 i について arr[i] += mat[i][q] を行い、「列 p〜q を使ったときの行ごとの合計」を表す1次元配列 arr を更新します。
   • ハッシュマップ pcnt を {0: 1} で初期化し、累積和 pref を 0 にします。
   • i = 0 から n-1 まで順に pref += arr[i] で累積和を更新し、pref - sumTarget が pcnt に登録済みなら、その出現回数を ans に加算します。
   • 最後に pcnt[pref]++ で現在の累積和の出現回数を記録します。
5. すべての列ペアの走査が完了したら ans を返します。

ポイントは、列の範囲 [p, q] を固定すると「合計が target になる行方向の連続区間」の数え上げが1次元問題に帰着できる点です。ハッシュマップによって必要な累積和の値を O(1) で参照できるため、行列全体を高速に処理できます。

C++の実装例

以下は、上記のアルゴリズムを実装したC++のコードです。

#include<bits/stdc++.h>

using namespace std;

int solve(vector<vector<int>>& mat, int sumTarget) {
    int n = mat.size();
    int m = n == 0 ? 0 : mat[0].size();

    if (m > n) {
        vector<vector<int>> transpose(m, vector<int>(n));
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < m; j++) {
                transpose[j][i] = mat[i][j];
            }
        }
        return solve(transpose, sumTarget);
    }

    int ans = 0;
    for (int p = 0; p < m; p++) {
        vector<int> arr(n);
        for (int q = p; q < m; q++) {
            for (int i = 0; i < n; i++) {
                arr[i] += mat[i][q];
            }

            unordered_map<int, int> pcnt = {{0, 1}};
            int pref = 0;
            for (int i = 0; i < n; i++) {
                pref += arr[i];
                auto tmp = pcnt.find(pref - sumTarget);
                if (tmp != end(pcnt)) ans += tmp->second;
                pcnt[pref]++;
            }
        }
    }
    return ans;
}

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, 5) << endl;
    return 0;
}

入力

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

出力

3

計算量の目安

• 時間計算量: O(min(m, n)² × max(m, n)) — 列ペアの組み合わせが最大 min(m, n)² 通りあり、各ペアで1次元の走査に max(m, n) かかるためです。
• 空間計算量: O(max(m, n)) — 行ごとの合計を保持する配列 arr とハッシュマップ pcnt に必要な分です。