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

C++로 푸는 부분 배열 최솟값의 합: 단조 스택 알고리즘 완벽 가이드

문제 설명

정수 배열 A가 주어졌을 때, A에서 만들 수 있는 모든 연속된(contiguous) 부분 배열 B에 대해 min(B)의 값을 구하고 그 합을 계산하는 문제입니다. 답이 매우 커질 수 있으므로 결과는 109 + 7로 나눈 나머지(modulo) 형태로 반환해야 합니다.

예를 들어 입력 배열이 [3,1,2,4]라고 가정해 보겠습니다. 만들 수 있는 부분 배열은 [3], [1], [2], [4], [3,1], [1,2], [2,4], [3,1,2], [1,2,4], [3,1,2,4]로 총 10개이며, 각 부분 배열의 최솟값은 순서대로 [3,1,2,4,1,1,2,1,1,1]입니다. 이 값들을 모두 더하면 17이 되므로 정답은 17입니다.

접근 방법: 단조 스택(Monotonic Stack) 활용

모든 부분 배열을 직접 생성하면 O(n²)개가 나오고, 각각의 최솟값을 찾는 데 추가 비용까지 들기 때문에 비효율적입니다. 대신 "각 원소가 최솟값이 되는 부분 배열은 몇 개인가?"를 원소별로 계산하면 전체 합을 한 번에 구할 수 있습니다. 이때 단조 스택을 사용하면 각 원소의 좌우 경계를 O(n) 시간에 효율적으로 찾을 수 있습니다.

알고리즘 단계

  • 모듈로 상수 m := 109 + 7로 설정합니다.
  • 오버플로를 방지하기 위해 두 개의 헬퍼 메서드를 정의합니다. add(a, b)는 (a mod m + b mod m) mod m을, mul(a, b)는 (a mod m × b mod m) mod m을 반환합니다.
  • 메인 메서드에서는 배열 A를 받아 스택 st를 선언하고, n := 배열 A의 크기로 설정합니다.
  • 크기가 n인 배열 left를 -1로, 같은 크기의 배열 right를 n으로 초기화합니다. left[i]는 i의 왼쪽에서 A[i]보다 작은 값 중 가장 가까운 인덱스를, right[i]는 오른쪽 방향에서 같은 조건을 만족하는 인덱스를 저장합니다.
  • ans := 0으로 초기화합니다.
  • i를 0부터 n−1까지 순회하며 다음을 수행합니다.
    • 스택이 비어 있지 않고 A[스택 top] ≥ A[i]인 동안 pop합니다.
    • 스택이 비어 있지 않으면 left[i] := 스택의 top으로 설정합니다.
    • i를 스택에 push합니다.
  • 스택을 모두 비웁니다.
  • 이번에는 i를 n−1부터 0까지 역순으로 순회하며 동일한 과정을 수행해 right 배열을 채웁니다.
  • 마지막으로 i를 0부터 n−1까지 순회하며 다음을 수행합니다.
    • leftBound := i − (left[i] + 1), rightBound := (right[i] − 1) − i로 설정합니다.
    • contri := 1 + leftBound + rightBound + (leftBound × rightBound). 이 값은 A[i]가 최솟값이 되는 부분 배열의 개수, 즉 해당 원소의 기여도(contribution)입니다.
    • ans := add(ans, mul(contri, A[i]))로 누적합니다.
  • ans를 반환합니다.

참고로 왼쪽 패스에서는 '≥' 조건을, 오른쪽 패스에서는 '>' 조건을 사용하는데, 이러한 비대칭 처리 덕분에 중복된 값이 배열에 포함되어 있어도 각 부분 배열이 정확히 한 번씩만 계산됩니다.

C++ 구현 예제

아래 구현을 통해 더 잘 이해해 보겠습니다.

#include <bits/stdc++.h>
using namespace std;
typedef long long int lli;
const lli MOD = 1e9 + 7;
class Solution {
public:
   lli add(lli a, lli b){
      return (a % MOD + b % MOD) % MOD;
   }
   lli mul(lli a, lli b){
      return (a % MOD * b % MOD) % MOD;
   }
   int sumSubarrayMins(vector<int>& A) {
      stack <int> st;
      int n = A.size();
      vector <int> left(n, -1);
      vector <int> right(n, n);
      int ans = 0;
      for(int i = 0; i < n; i++){
         while(!st.empty() && A[st.top()] >= A[i]){
         st.pop();
      }
      if(!st.empty())left[i] = st.top();
         st.push(i);
      }
      while(!st.empty())st.pop();
      for(int i = n - 1; i >= 0; i--){
         while(!st.empty() && A[st.top()] > A[i]){
            st.pop();
         }
         if(!st.empty())right[i] = st.top();
            st.push(i);
      }
      for(int i = 0; i < n; i++){
         int leftBound = i - (left[i] + 1);
         int rightBound = (right[i] - 1) - i;
         int contri = 1 + leftBound + rightBound + (leftBound * rightBound);
         ans = add(ans, mul(contri, A[i]));
      }
      return ans;
   }
};
main(){
   vector<int> v = {3,1,2,4};
   Solution ob;
   cout << (ob.sumSubarrayMins(v));
}

입력

[3,1,2,4]

출력

17

복잡도 분석

시간 복잡도는 배열을 세 번 순회하지만 각 원소가 스택에 최대 한 번 push되고 한 번 pop되므로 O(n)입니다. 공간 복잡도는 left, right 배열과 스택 때문에 O(n)입니다. 브루트포스 방식의 O(n²) 또는 O(n³)에 비해 상당히 효율적인 접근법입니다.