누적 합(Prefix Sum) 개념과 구간 합 구하기

[파이썬 알고리즘] 누적 합(Prefix Sum) 개념과 '구간 합 구하기 4' 문제 풀이

1. 문제 이해: 왜 누적 합이 필요할까?

'구간 합 구하기 4' 문제는 N개의 숫자가 주어졌을 때, 특정 구간 [i, j]의 합을 M번 구하는 문제입니다.

입력 조건:

• N과 M은 최대 100,000
• 각 구간의 합을 M번 계산해야 함

만약 단순하게 매번 i부터 j까지 반복문을 돌려 합을 구한다면, 한 번의 계산에 O(N)의 시간이 걸립니다. 이를 M번 반복하면 총 O(M x N)의 시간 복잡도가 되어, 100,000 x 100,000 = 10의 10승이라는 엄청난 연산 횟수로 시간 초과가 발생하게 됩니다.

이 비효율을 해결해 줄 열쇠가 바로 누적 합(Prefix Sum)입니다.

2. 핵심 개념: 누적 합 배열(Prefix Sum Array) 만들기

누적 합 배열을 이해하는 가장 좋은 방법은 저금통을 상상하는 것입니다. 매일 저금통에 돈을 넣을 때, 오늘까지 넣은 총액을 매번 기록한다고 생각해 보세요.

요일	오늘 넣은 돈	오늘까지 총액
월요일	500원	500원
화요일	400원	900원 (500 + 400)
수요일	300원	1,200원 (900 + 300)
목요일	200원	1,400원 (1,200 + 200)
금요일	100원	1,500원 (1,400 + 100)

여기서 '오늘까지 총액'이 바로 누적 합입니다. 누적 합 배열(Prefix Sum Array)은 이처럼 원본 배열의 첫 번째 원소부터 현재 위치까지의 합을 미리 계산해 저장해 둔 배열입니다. 이 배열을 만드는 데는 원본 배열을 딱 한 번만 순회하면 됩니다.

누적 합 배열 생성 과정

원본 배열 arr가 [5, 4, 3, 2, 1]일 때, sum_arr라는 누적 합 배열을 만들어볼게요. 계산의 편의를 위해 sum_arr[0]은 0으로 초기화하고 시작하는 것이 일반적입니다.

• sum_arr[1] = sum_arr[0] + arr[0] = 0 + 5 = 5
• sum_arr[2] = sum_arr[1] + arr[1] = 5 + 4 = 9
• sum_arr[3] = sum_arr[2] + arr[2] = 9 + 3 = 12
• sum_arr[4] = sum_arr[3] + arr[3] = 12 + 2 = 14
• sum_arr[5] = sum_arr[4] + arr[4] = 14 + 1 = 15

결과적으로 sum_arr는 [0, 5, 9, 12, 14, 15]가 됩니다.

이 과정을 코드로 표현하면 다음과 같습니다.

arr = [5, 4, 3, 2, 1]
n = len(arr)
sum_arr = [0] * (n + 1)

for i in range(1, n + 1):
    sum_arr[i] = sum_arr[i-1] + arr[i-1]

이렇게 O(N)의 시간 복잡도로 누적 합 배열을 미리 만들어두면, 이제 어떤 구간의 합이든 단 한 번의 연산만으로 아주 빠르게 구할 수 있습니다. 이것이 바로 누적 합 알고리즘의 핵심적인 강점입니다.

3. 구간 합을 O(1)에 구하는 방법

누적 합 배열을 만들었다면, 이제 어떤 구간의 합이든 단 한 번의 연산으로 구할 수 있습니다.

[i, j] 구간의 합 = sum_arr[j] - sum_arr[i-1]

이 공식이 성립하는 이유는 다음과 같습니다.

• sum_arr[j]는 첫 번째 원소부터 j번째 원소까지의 총합입니다.
• sum_arr[i-1]은 첫 번째 원소부터 i-1번째 원소까지의 총합입니다.
• 따라서 sum_arr[j]에서 sum_arr[i-1]을 빼면, i번째부터 j번째까지의 합만 남게 됩니다.

예시: [2, 4] 구간의 합

• sum_arr[4] = 14 (1번째 ~ 4번째 원소의 합)
• sum_arr[1] = 5 (1번째 ~ 1번째 원소의 합)
• sum_arr[4] - sum_arr[1] = 14 - 5 = 9
• 이는 arr의 4 + 3 + 2 = 9와 일치합니다.

4. 파이썬 코드 구현 (solution1)

이 원리를 바탕으로 파이썬 코드를 작성해 봅시다. sys.stdin.readline을 사용해 입력 속도를 최적화하는 것이 중요합니다.

import sys

def solve():
    # 1. 입력 처리
    input = sys.stdin.readline
    n, m = map(int, input().split())
    arr = list(map(int, input().split()))

    # 2. 누적 합 배열 생성 (O(N))
    # 계산 편의를 위해 sum_arr[0]을 0으로 설정
    sum_arr = [0] * (n + 1)
    for i in range(1, n + 1):
        sum_arr[i] = sum_arr[i-1] + arr[i-1]

    # 3. M번의 구간 합 계산 및 출력 (O(M))
    for _ in range(m):
        i, j = map(int, input().split())
        print(sum_arr[j] - sum_arr[i-1])

if __name__ == "__main__":
    solve()

이 코드는 합 배열을 만드는 데 O(N)의 시간, 각 구간 합을 구하는 데 O(1)의 시간을 사용하므로, 총 O(N+M의 시간 복잡도로 문제를 해결합니다.

보너스: 더 간결한 코드 (solution2)

파이썬의 itertools.accumulate 함수를 사용하면 누적 합 배열을 더 간결하게 만들 수 있습니다.

import sys
from itertools import accumulate

def solve_with_accumulate():
    input = sys.stdin.readline
    n, m = map(int, input().split())
    numbers = list(map(int, input().split()))

    # itertools.accumulate 활용 (누적 합 계산)
    prefix_sum = [0] + list(accumulate(numbers))

    for _ in range(m):
        i, j = map(int, input().split())
        print(prefix_sum[j] - prefix_sum[i-1])

if __name__ == "__main__":
    solve_with_accumulate()

accumulate 함수는 for 루프를 대체하여 코드를 더욱 파이써닉(Pythonic)하게 만들어줍니다.

5. 마무리

누적 합 알고리즘은 '구간 합 구하기 4'와 같은 문제의 핵심 풀이법이며, 다양한 응용 문제에도 적용될 수 있습니다. 이 개념을 확실히 이해하고 연습한다면 코딩 테스트에서 만나는 유사 문제들을 자신 있게 해결할 수 있을 것입니다.

다음에는 또 다른 유용한 알고리즘으로 찾아뵙겠습니다. 감사합니다!