빠르게 배열의 부분 합을 구하는 Prefix Sum 알고리즘

Prefix Sum 알고리즘이란?

문제를 풀다 보면 배열에서 특정 구간의 합을 구해야 하는 문제들을 종종 만나곤 한다. 한 번만 계산해야 한다면 반복문을 통해 \(O(N)\)의 시간 복잡도로 해결할 수 있지만, 여러 번 반복해서 계산해야 한다면 조금 더 효율적인 알고리즘이 필요하다. 이때 사용할 수 있는 알고리즘이 Prefix Sum이다.

 

Prefix Sum 알고리즘은 배열의 첫 번째 원소부터 각 원소들까지의 합을 미리 계산해두는 알고리즘이다. 예를 들어, 배열 [1, 2, 3, 4, 5]가 있을 때 첫 번째 원소의 Prefix Sum은 1, 두 번째 원소의 Prefix Sum은 1 + 2 = 3, 세 번째 원소의 Prefix Sum은 1 + 2 + 3 = 6 이 된다. 이렇게 미리 배열의 모든 원소들의 Prefix Sum을 구해 놓으면 배열의 특정 범위 내 합을 빠르게 구할 수 있다.

 

1차원 배열에서 구간 합 빠르게 구하기

Prefix Sum 배열 계산하기

1차원 배열에서 Prefix Sum은 \(psum[i] = psum[i - 1] + arr[i]\) 공식을 이용해 계산할 수 있다. 

vector<int> getPartialSumVector(const vector<int>& arr) {
    vector<int> psum(arr.size());
    
    psum[0] = arr[0];
    
    for (int i = 1 ; i < arr.size() ; ++i) {
        psum[i] = psum[i - 1] + arr[i];
    }
    
    return psum;
}

 

구간 합 구하기

이렇게 Prefix Sum 배열을 만들고 나면 원본 배열의 특정 범위 \([l, r]\)의 합을 다음과 같이 계산할 수 있다.

$$sum(l, r) = psum[r] - psum[l - 1]$$

\(psum[r]\)은 배열의 첫 번째 원소부터 \(r\)번째 원소까지의 합을 나타낸다. \(psum[l - 1]\)은 배열의 첫 번째 원소부터 \(l - 1\)번째 원소까지의 합을 나타낸다. 따라서 \(psum[r]\)에서 \(psum[l - 1]\)을 빼면, 배열의 \(l\)번째 원소부터 \(r\)번째 원소의 합만 남게 된다.

int getSumOf(const vector<int>& psum, int l, int r) {
    if (l == 0) {
        return psum[r];
    }
    
    return psum[r] - psum[l - 1];
}

\(O(N)\)의 시간 복잡도를 이용해 Prefix Sum 배열을 만들고 나면 이후 배열의 구간 합을 구하는 연산은 \(O(1)\)에 할 수 있다.

 

2차원 배열에서 구간 합 빠르게 구하기

Prefix Sum 배열 계산하기

2차원 배열의 Prefix Sum의 정의는 1차원 배열의 Prefix Sum 정의와는 조금 다르다. 2차원 배열에서  \(y_{ij}\)의 Prefix Sum은 (0, 0)에서 (i, j)까지의 합이라고 정의한다. 쉽게 말해서 2차원 배열의 (0, 0)부터 (i, j)까지 사각형을 그린 뒤 사각형 안에 포함되는 원소들을 다 합친것을 의미한다. 2차원 배열의 Prefix Sum 배열은 \(psum[i][j] = arr[i][j] + psum[i - 1][j] + psum[i][j - 1] - psum[i - 1][j - 1]\) 공식을 통해 계산할 수 있다.

2차원 배열의 Prefix Sum 배열

수식으로만 보면 다소 헷갈릴 수 있어 그림과 함께 설명을 다시 한다. (i, j)의 Prefix Sum을 계산할 때 이전에 계산한 Prefix Sum 값을 재활용한다. 다음과 같은 영역을 생각해보자.

• 분홍색 영역: (0, 0)부터 (i, j - 1)까지의 합
• 초록색 영역: (0, 0)부터 (i - 1, j)까지의 합
• 검은색 영역: 분홍색 영역과 초록색 영역이 겹치는 부분

(i, j)의 Prefix Sum을 계산하려면 분홍색 영역과 초록색 영역을 더한 뒤, 중복해서 더해진 검은색 영역을 빼야 한다. 마지막으로 (i, j)의 값을 더해주면 (0, 0)부터 (i, j)까지의 Prefix Sum을 계산할 수 있다.

vector<vector<int>> getPartialSumVector(const vector<vector<int>>& arr) {
    int rows = arr.size();
    int cols = arr[0].size();
    
    vector<vector<int>> psum(rows, vector<int>(cols));
    
    for (int i = 0 ; i < rows ; ++i) {
        for (int j = 0 ; j < cols ; ++j) {
            psum[i][j] = arr[i][j];
            
            if (i > 0) {
                psum[i][j] += psum[i - 1][j];
            }
            
            if (j > 0) {
                psum[i][j] += psum[i][j - 1];
            }
            
            if (i > 0 && j > 0) {
                psum[i][j] -= psum[i - 1][j - 1];
            }
        }
    }
    
    return psum;
}

 

구간 합 구하기

Prefix Sum 배열을 계산할 때와 비슷한 방법을 이용해 2차원 배열에서 (l1, r1) ~ (l2, r2)의 구간 합을 구할 수 있다. 이를 수식으로 표현하면 다음과 같다.

$$sum(l1, r1, l2, r2) = psum[l2][r2] - psum[l1 - 1][r2] - psum[l2][r1 - 1] + psum[l1 - 1][r1 - 1]$$

2차원 배열의 구간 합

구간 합을 구하는 과정 역시 수식으로만 보면 헷갈릴 수 있어 그림과 함께 다시 봐보자. (l1, r1) ~ (l2, r2)의 구간 합을 구하는 과정에 사용할 Prefix Sum을 다음과 같이 생각해 보자

• 빨간색 영역: (0, 0)에서 (l2, r2)까지의 합
• 초록색 영역: (0, 0)에서 (l2, r1 - 1)까지의 합
• 파란색 영역: (0, 0)에서 (l1 - 1, r2)까지의 합
• 검은색 영역: 초록색 영역과 파란색 영역이 겹치는 부분

구간 (l1, r1) ~ (l2, r2)의 합을 구하기 위해서는 빨간색 영역에서 초록색 영역과 파란색 영역을 뺀 뒤, 두 번 빠진 검은색 영역을 한 번 더해줘야 한다.

int getSumOf(const vector<vector<int>>& psum, int l1, int r1, int l2, int r2) {
    int ret = psum[l2][r2];
    
    if (l1 > 0) {
        ret -= psum[l1 - 1][r2];
    }
    
    if (r1 > 0) {
        ret -= psum[l2][r1 - 1];
    }
    
    if (l1 > 0 && r1 > 0) {
        ret += psum[l1 - 1][r1 - 1];
    }
    
    return ret;
}

\(O(NM)\)의 시간 복잡도를 이용해 2차원 배열의 Prefix Sum 배열을 만들고 나면 그다음 특정 구간의 합을 구하는 연산은 \(O(1)\)에 할 수 있다.

 

Prefix Sum을 이용해 풀 수 있는 문제들

• [백준 2167번] 2차원 배열의 합
• [백준 11659번] 구간 합 구하기 4
• [백준 2559번] 수열