C++로 세 개의 겹치지 않는 부분 배열의 최대 합 구하기

문제 소개

양의 정수로 이루어진 배열 nums가 주어졌을 때, 합이 최대가 되는 세 개의 서로 겹치지 않는(non-overlapping) 부분 배열을 찾아야 합니다. 각 부분 배열의 길이는 k로 고정되어 있으며, 세 부분 배열에 속한 모든 원소(총 3×k개)의 합을 최대화하는 것이 목표입니다.

결과는 각 구간의 시작 인덱스를 담은 리스트 형태로 반환합니다. 만약 조건을 만족하는 답이 여러 개라면, 그중 사전순으로 가장 작은(lexicographically smallest) 하나를 선택해야 합니다.

예시로 이해하기

입력 배열이 [1,2,1,2,6,8,4,1]이고 k = 2라고 가정해 보겠습니다. 이 경우 정답은 [0, 3, 5]입니다. 즉, 시작 인덱스 0, 3, 5에 해당하는 부분 배열은 각각 [1,2], [2,6], [8,4]이며, 이 조합이 만들 수 있는 합 중 가장 큰 값을 가집니다.

해결 접근 방식

이 문제는 누적합(prefix sum)과 사전 계산(precomputation) 기법을 활용하면 O(n) 시간 복잡도로 효율적으로 해결할 수 있습니다. 전체 흐름은 다음과 같습니다.

n: 배열 nums의 크기를 저장합니다.
ret: 결과를 담을 크기 3의 배열을 선언합니다.
sum: 크기 n+1의 누적합 배열을 만들어, 임의의 구간 합을 O(1)에 구할 수 있도록 준비합니다.
posLeft: 각 인덱스 i를 기준으로, 왼쪽 영역에서 길이 k짜리 부분 배열의 합이 최대가 되는 시작 위치를 저장합니다. 동점일 때는 먼저 발견된 더 작은 인덱스를 유지하여 사전순 조건을 만족합니다.
posRight: 각 인덱스 i를 기준으로, 오른쪽 영역에서 길이 k짜리 부분 배열의 합이 최대가 되는 시작 위치를 저장합니다. 동점일 때 더 작은 인덱스를 선택할 수 있도록 '>=' 비교를 사용합니다.
마지막으로 가운데 부분 배열의 시작 위치 i(k ≤ i ≤ n−2k)를 순회하면서, posLeft[i−1], i, posRight[i+k] 세 위치의 부분 배열 합을 계산하고 최댓값을 갱신합니다.

핵심 아이디어는 '왼쪽 최적 위치'와 '오른쪽 최적 위치'를 미리 계산해 두면, 가운데 위치 하나만 결정했을 때 나머지 두 위치를 O(1)에 얻을 수 있다는 점입니다. 덕분에 전체 탐색을 O(n²)이 아닌 O(n)에 완료할 수 있습니다.

C++ 구현 코드

#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<auto> v){
    cout << "[";
    for(int i = 0; i<v.size(); i++){
        cout << v[i] << ", ";
    }
    cout << "]"<<endl;
}
class Solution {
public:
    vector<int> maxSumOfThreeSubarrays(vector<int>& nums, int k) {
        int n = nums.size();
        vector <int> ret(3, INT_MAX);
        vector <int> sum(n + 1);
        for(int i = 0; i < n; i++){
            sum[i + 1] = sum[i] + nums[i];
        }
        vector <int> posLeft(n);
        vector <int> posRight(n, n - k);
        for(int i = k, currMax = sum[k] - sum[0]; i < n; i++){
            int newTotal = sum[i + 1] - sum[i + 1- k];
            if(newTotal > currMax){
                currMax = newTotal;
                posLeft[i] = i + 1 - k;
            }else{
                posLeft[i] = posLeft[i - 1];
            }
        }
        for(int i = n - k - 1, currMax = sum[n] - sum[n - k]; i >=0 ; i--){
            int newTotal = sum[i + k] - sum[i];
            if(newTotal >= currMax){
                currMax = newTotal;
                posRight[i] = i;
            }else{
                posRight[i] = posRight[i + 1];
            }
        }
        int req = 0;
        for(int i = k; i <= n - 2 * k; i++){
            int l = posLeft[i - 1];
            int r = posRight[i + k];
            int temp = (sum[l + k] - sum[l]) + (sum[i + k] - sum[i]) + (sum[r + k] - sum[r]);
            if(temp > req){
                ret[0] = l;
                ret[1] = i;
                ret[2] = r;
                req = temp;
            }
        }
        return ret;
    }
};
main(){
    Solution ob;
    vector<int> v = {1,2,1,2,6,8,4,1};
    print_vector(ob.maxSumOfThreeSubarrays(v, 2));
}

실행 결과 확인

위 코드를 앞서 살펴본 예시 입력으로 실행하면 다음과 같은 결과를 얻을 수 있습니다.

입력

{1,2,1,2,6,8,4,1}
2

출력

[0, 3, 5]

출력 결과는 예상대로 시작 인덱스 [0, 3, 5]를 반환하며, 해당 부분 배열들의 합 24로 세 개의 비중첩 부분 배열 중 최대 합을 달성함을 확인할 수 있습니다.