C++でターゲットの合計値になる部分行列の数を求める方法

問題概要

行列(matrix)とターゲット値(target)が与えられたとき、要素の合計がターゲットと一致する空でない部分行列がいくつ存在するかを求める問題です。ここで部分行列 [(x1, y1), (x2, y2)] とは、x が x1 以上 x2 以下、y が y1 以上 y2 以下の範囲に含まれるすべてのセル matrix[x][y] の集合を指します。また、2つの部分行列 [(x1, y1), (x2, y2)] と [(x1', y1'), (x2', y2')] は、座標が1つでも異なれば別のものとして扱われます(例えば x1 ≠ x1' の場合など)。

入力例

このとき target = 0 とすると、出力は 4 になります。これは、0 のみを含む 1×1 の部分行列が4つ存在するためです。

アルゴリズムの手順

この問題は、累積和(プレフィックスサム)とハッシュマップを組み合わせることで効率的に解くことができます。具体的な手順は以下の通りです。

• ans := 0(答えを格納する変数)

• col := 列数、row := 行数

• 各行 i について、j = 1 から col-1 まで matrix[i][j] += matrix[i][j-1] と更新し、行ごとの累積和を作成する

• ハッシュマップ m を定義する

• i = 0 から col-1 までループ:

  • j = i から col-1 までループ:

    • マップ m をクリアし、m[0] := 1、sum := 0 で初期化する

    • k = 0 から row-1 まで各行について:

      • current := matrix[k][j](i - 1 >= 0 の場合は current -= matrix[k][i-1])

      • sum := sum + current

      • ans := ans + m[target - sum]

      • m[-sum] を 1 増やす

• ans を返す

アルゴリズムのポイント

この手法は、1次元配列における「合計が K になる部分配列の個数を数える」問題を2次元へ拡張したものです。列のペア (i, j) を固定すると、各行におけるその列範囲の値は1次元配列とみなせます。そこで累積和 sum を管理しながらハッシュマップを参照することで、「現在の累積和から target を引いた値」と等しい過去の累積和が何回現れたかを即座に取得でき、条件を満たす部分行列を効率的に数え上げられます。計算量は O(col² × row) となり、全部分行列を素朴に走査する方法より大幅に高速です。

C++実装例

#include <bits/stdc++.h>
using namespace std;
class Solution {
   public:
   int numSubmatrixSumTarget(vector<vector<int>>& matrix, int
   target) {
      int ans = 0;
      int col = matrix[0].size();
      int row = matrix.size();
      for(int i = 0; i < row; i++){
         for(int j = 1; j < col; j++){
            matrix[i][j] += matrix[i][j - 1];
         }
      }
      unordered_map <int, int> m;
      for(int i = 0; i < col; i++){
         for(int j = i; j < col; j++){
            m.clear();
            m[0] = 1;
            int sum = 0;
            for(int k = 0; k < row; k++){
               int current = matrix[k][j];
               if(i - 1 >= 0)current -= matrix[k][i - 1];
               sum += current;
               ans += m[target - sum];
               m[-sum]++;
            }
         }
      }
      return ans;
   }
};
main(){
   Solution ob;
   vector<vector<int>> v = {{0,1,0},{1,1,1},{0,1,0}};
   cout << (ob.numSubmatrixSumTarget(v, 0));
}

入力

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

出力

4