C++에서 모든 원소가 1인 부분 행렬의 개수 구하기

문제 소개

1과 0으로만 이루어진 행렬이 주어질 때, 행렬 안에서 모든 원소가 1로만 구성된 부분 행렬(submatrix)의 개수를 구하는 프로그램을 만들어 보겠습니다.

예를 들어 아래와 같은 행렬이 입력으로 주어진다고 가정해 보겠습니다.

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

이 경우 정답은 12입니다.

접근 방법

이 문제는 2차원 누적 합(prefix sum) 배열과 이진 탐색(binary search)을 조합하면 효율적으로 해결할 수 있습니다. 핵심 아이디어는 다음과 같습니다.

누적 합 배열을 미리 계산해 두면, 임의의 직사각형 영역의 합을 O(1) 시간에 구할 수 있습니다.
값이 1인 각 칸을 왼쪽 위 꼭짓점으로 하는 부분 행렬을 차례대로 고려합니다.
높이 k를 하나씩 늘려 가며, 이진 탐색으로 해당 높이에서 만들 수 있는 '모든 원소가 1'인 최대 너비를 구합니다.
찾은 최대 너비만큼 부분 행렬 개수를 정답에 누적합니다.

알고리즘 단계

n := 행렬의 행 개수, m := 행렬의 열 개수로 설정합니다.
(n+1) × (m+1) 크기의 2차원 배열 add(누적 합 배열)를 정의하고 0으로 초기화합니다.
모든 칸 (i, 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]
res := 0으로 초기화합니다.
각 칸 (i, j)에 대해 다음을 수행합니다.
- matrix[i][j]가 0이면 건너뛰고 다음 반복으로 넘어갑니다.
- 높이 k를 1부터 (n - i)까지 하나씩 늘리며 다음을 반복합니다.
  - p := 0, q := m - j로 설정한 뒤 이진 탐색을 시작합니다.
  - p ≤ q인 동안 다음을 반복합니다.
    - x := (p + q) / 2
    - a := k * x (높이 k, 너비 x인 영역의 이상적인 합)
    - cur := add[i + k][j + x] - add[i][j + x] - add[i + k][j] + add[i][j]
    - cur이 a와 같다면 해당 영역은 모두 1이므로 r := x로 저장하고, p := x + 1로 더 넓은 폭을 탐색합니다.
    - 그렇지 않다면 q := x - 1로 탐색 범위를 줄입니다.
  - r이 0이면 더 이상 확장할 수 없으므로 내부 반복을 종료합니다.
  - 그렇지 않으면 res := res + r로 정답을 누적합니다.
모든 반복이 끝나면 res를 반환합니다.

동작 원리

여기서 cur은 시작점 (i, j)에서 오른쪽으로 x칸, 아래로 k칸 뻗은 직사각형 영역의 실제 합입니다. 영역이 모두 1이라면 cur은 k * x와 정확히 일치합니다. 따라서 이진 탐색으로 '모두 1'인 최대 너비 r을 빠르게 찾을 수 있으며, 높이 k에서 만들 수 있는 부분 행렬의 개수는 정확히 r개가 됩니다.

예제 코드 (C++)

아래는 위 알고리즘을 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}}

출력

시간 복잡도

각 칸마다 최대 n개의 높이에 대해 O(log m)의 이진 탐색을 수행하므로, 전체 시간 복잡도는 O(n² · m · log m) 수준입니다. 모든 부분 행렬을 일일이 검사하는 브루트 포스 방식에 비해 훨씬 효율적으로 동작합니다.