Computer >> 컴퓨터 >  >> 프로그래밍 >> C++

C++로 주어진 행렬에서 1로만 이루어진 부분 행렬의 개수 계산하기

문제 개요

2차원 이진 행렬(0과 1로만 구성된 행렬)이 주어졌을 때, 모든 요소가 1인 부분 행렬(submatrix)의 총 개수를 구하는 것이 이번 문제의 목표입니다.

예를 들어 입력 행렬이 다음과 같다고 가정해 보겠습니다.

110
110
001

이때 출력은 10이 됩니다. 그 이유를 살펴보면, 1×1 행렬 5개, 2×1 행렬 2개, 1×2 행렬 2개, 그리고 2×2 행렬 1개가 존재하기 때문입니다.

접근 방법

이 알고리즘의 핵심 아이디어는 각 행을 히스토그램의 바닥으로 보는 것입니다. 배열 temp의 j번째 원소는 '현재 행에서 위 방향으로 연속된 1의 개수', 즉 열 j에서 현재 행을 밑변으로 하는 히스토그램의 높이를 의미합니다. 여기에 단조 스택(monotonic stack)을 활용하면 현재 행을 아래 변으로 하는 부분 행렬의 개수를 선형 시간에 효율적으로 셀 수 있습니다.

알고리즘 단계

이 문제를 해결하기 위해 다음 단계를 따릅니다.

  • 배열 a를 매개변수로 받는 getAns() 함수를 정의합니다.
  • ret을 0으로 초기화하고, n에 배열 a의 크기를 저장합니다.
  • 크기가 n인 배열 v와 하나의 스택 st를 정의합니다.
  • i를 0부터 배열 a의 크기 미만까지 1씩 증가시키며 다음을 반복합니다.
    • 스택 st가 비어 있지 않고 a[st.top()] ≥ a[i]인 동안 st에서 요소를 꺼냅니다(pop).
    • 스택이 비어 있지 않다면:
      • prev에 스택 최상단 값을 저장합니다.
      • v[i]에 v[prev]를 더합니다.
      • v[i]에 a[i] × (i − prev)를 더합니다.
    • 그렇지 않다면:
      • v[i]에 a[i] × (i + 1)을 더합니다.
    • i를 스택 st에 삽입합니다.
  • 배열 v의 모든 요소를 순회하며 ret에 누적한 뒤, ret을 반환합니다.

메인 solve() 메서드에서는 다음을 수행합니다.

  • ret을 0으로 초기화하고, n에 행렬 v의 행 개수를 저장합니다.
  • m에는 (n이 0이 아니면 v[0]의 크기, 그렇지 않으면 0)을 저장합니다.
  • 크기가 m인 배열 temp를 정의합니다.
  • i를 0부터 n 미만까지 1씩 증가시키며 반복합니다.
    • j를 0부터 m 미만까지 1씩 증가시키며 반복하면서, v[i][j]가 0이 아니면 temp[j]를 temp[j] + 1로, 0이면 0으로 갱신합니다.
    • ret에 getAns(temp)의 반환값을 더합니다.
  • 최종적으로 ret을 반환합니다.

예제 코드

더 나은 이해를 위해 다음 C++ 구현 예제를 살펴보겠습니다.

#include <bits/stdc++.h>
using namespace std;
class Solution {
   public:
   int getAns(vector<int>& a) {
      int ret = 0;
      int n = a.size();
      vector<int> v(n);
      stack<int> st;
      for (int i = 0; i < a.size(); i++) {
         while (!st.empty() && a[st.top()] >= a[i])
            st.pop();
         if(!st.empty()) {
            int prev = st.top();
            v[i] += v[prev];
            v[i] += a[i] * (i - prev);
         }
         else{
            v[i] += a[i] * (i + 1);
         }
         st.push(i);
      }
      for (int i : v) {
         ret += i;
      }
      return ret;
   }
   int solve(vector<vector<int>>& v) {
      int ret = 0;
      int n = v.size();
      int m = n ? v[0].size() : 0;
      vector<int> temp(m);
      for (int i = 0; i < n; i++) {
         for (int j = 0; j < m; j++) {
            temp[j] = v[i][j] ? temp[j] + 1 : 0;
         }
         ret += getAns(temp);
      }
      return ret;
   }
};
int solve(vector<vector<int>>& matrix) {
   return (new Solution())->solve(matrix);
}
int main(){
   vector<vector<int>> matrix = {
      {1, 1, 0},
      {1, 1, 0},
      {0, 0, 1}
   };
   cout << solve(matrix);
}

입력

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

출력

10

복잡도 분석

각 행마다 단조 스택 기반의 getAns()가 한 번씩 호출되므로, 전체 시간 복잡도는 O(n × m)입니다. 여기서 n은 행의 개수, m은 열의 개수입니다. 공간 복잡도는 히스토그램 높이 배열과 스택에 사용되는 O(m)입니다. 완전 탐색으로 모든 부분 행렬을 검사하는 O(n² × m²) 방식에 비해 훨씬 효율적이라는 점이 이 접근법의 가장 큰 장점입니다.