C++로 요소의 합이 특정 값과 같은 부분행렬 개수 구하기

정수 값으로 구성된 행렬이 주어졌을 때, 요소들의 합이 특정 목표값(target)과 일치하는 부분행렬(submatrix)의 개수를 구하는 문제입니다. 여기서 부분행렬이란 원본 행렬에서 연속된 행과 열을 선택하여 만들 수 있는 직사각형 형태의 작은 행렬을 의미합니다.

문제 예시

다음과 같은 4×4 행렬이 있다고 가정해 보겠습니다.

0	0	1	0
0	1	0	0
0	1	0	1
1	1	0	1

목표값이 5일 때, 요소의 합이 5가 되는 부분행렬은 총 3개이므로 출력 결과는 3이 됩니다.

풀이 접근 방법

이 문제는 1차원 배열에서 '합이 k인 부분배열의 개수'를 세는 고전적인 기법, 즉 접두사 합(prefix sum)과 해시맵을 활용하는 방식을 행렬 문제로 확장한 것입니다.

두 개의 열 경계 p와 q를 고정한 뒤, 해당 열 범위에 포함되는 각 행의 요소 합을 1차원 배열에 누적합니다.
누적된 배열을 순회하며 접두사 합을 계산하고, (현재 접두사 합 − 목표값)이 해시맵에 등장한 횟수만큼 정답에 더해줍니다.
행렬의 열 개수가 행 개수보다 많으면(m > n), 탐색 효율을 높이기 위해 행렬을 전치(transpose)한 후 동일한 함수를 재귀적으로 호출합니다.

알고리즘 상세 단계

행렬의 행 개수를 n으로, 열 개수를 m으로 설정합니다. (n이 0이면 m도 0)
m이 n보다 크면 m×n 크기의 전치 행렬을 생성한 뒤 solve 함수를 재귀 호출합니다.
정답 변수 ans를 0으로 초기화합니다.
첫 번째 열 경계 p를 0부터 m−1까지 반복합니다.
- 길이가 n인 배열 arr를 선언합니다.
- 두 번째 열 경계 q를 p부터 m−1까지 반복하며, 각 행 i에 대해 arr[i] += mat[i][q]로 열 값을 누적합니다.
- 초깃값 {0, 1}을 가진 맵 pcnt와 접두사 합 변수 pref(=0)를 준비합니다.
- 행 인덱스 i를 0부터 n−1까지 순회하면서 pref += arr[i]로 접두사 합을 갱신하고, pcnt에서 (pref − sumTarget)을 조회합니다. 존재한다면 해당 값을 ans에 더한 후, pcnt[pref]를 1 증가시킵니다.
모든 반복이 끝나면 ans를 반환합니다.

이 방식의 시간 복잡도는 O(min(n,m)² × max(n,m))으로, 가능한 모든 부분행렬을 일일이 검사하는 O(n⁴) 완전 탐색에 비해 훨씬 효율적입니다.

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