C++ 범위 추가(Range Addition) 문제: 차분 배열로 O(n+k)에 해결하기

문제 개요

크기가 n인 배열이 주어지고, 모든 원소가 0으로 초기화되어 있다고 가정해 봅시다. 그리고 값 k가 함께 주어지며, 우리는 k번의 업데이트 연산을 수행해야 합니다.

각 연산은 [startIndex, endIndex, inc] 형태의 세 값으로 표현되며, 부분 배열 A[startIndex ... endIndex]의 startIndex부터 endIndex까지(양 끝 인덱스 포함) 모든 원소를 inc만큼 증가시킵니다. 목표는 k번의 모든 연산이 완료된 후의 최종 배열을 구하는 것입니다.

예를 들어 입력이 length = 5, updates = [[1,3,2],[2,4,3],[0,2,-2]]라면 출력은 [-2, 0, 3, 5, 3]이 됩니다.

단계별 해결 방법

이 문제는 다음 단계를 통해 해결할 수 있습니다.

크기 n의 결과 배열 ret을 정의합니다.
i := 0으로 초기화한 뒤, i가 배열 a의 크기보다 작은 동안 i를 1씩 증가시키며 다음을 반복합니다.
- l := a[i, 0]
- r := a[i, 1] + 1
- ret[l] := ret[l] + a[i, 2]
- 만약 r < n이라면 다음을 수행합니다.
  - ret[r] := ret[r] - a[i, 2]
i := 1로 초기화한 뒤, i가 n보다 작은 동안 i를 1씩 증가시키며 다음을 반복합니다.
- ret[i] := ret[i] + ret[i - 1]
배열 ret을 반환합니다.

핵심 아이디어: 차분 배열(Difference Array)

이 알고리즘의 핵심은 차분 배열 기법입니다. 매 업데이트마다 구간 내 모든 원소를 직접 변경하는 대신, 구간의 시작점에는 inc를 더하고 구간 끝의 다음 위치에는 inc를 빼서 변화량만 기록해 둡니다. 모든 업데이트를 처리한 후 한 번의 누적합(prefix sum) 계산만 수행하면 최종 배열을 얻을 수 있습니다.

이 방식을 사용하면 각 업데이트를 O(1) 시간에 처리할 수 있으므로 전체 시간 복잡도는 O(n + k)가 됩니다. 덕분에 구간이 길고 업데이트 횟수가 많은 경우에도 매우 효율적으로 동작합니다.

C++ 구현 예제

더 나은 이해를 돕기 위해 다음 구현 코드를 살펴보겠습니다.

#include <bits/stdc++.h>
using namespace std;

void print_vector(vector<int> v){
    cout << "[";
    for(int i = 0; i < v.size(); i++){
        cout << v[i] << ", ";
    }
    cout << "]" << endl;
}

class Solution {
public:
    vector<int> getModifiedArray(int n, vector<vector<int>>& a) {
        vector<int> ret(n);
        for (int i = 0; i < a.size(); i++) {
            int l = a[i][0];
            int r = a[i][1] + 1;
            ret[l] += a[i][2];
            if (r < n) {
                ret[r] -= a[i][2];
            }
        }
        for (int i = 1; i < n; i++) {
            ret[i] += ret[i - 1];
        }
        return ret;
    }
};

int main(){
    Solution ob;
    vector<vector<int>> v = {{1,3,2},{2,4,3},{0,2,-2}};
    print_vector(ob.getModifiedArray(5, v));
}

입력

5, {{1,3,2},{2,4,3},{0,2,-2}}

출력

[-2, 0, 3, 5, 3]