구간 합 구하기

언제 쓰나

“구간 [i, j]의 합을 여러 번 물어본다” — 질의가 반복되면 매번 더하는 O(n)이 누적되어 시간 초과
부분 배열의 합/평균, “합이 K인 구간”류 문제의 전처리
2차원이면 “부분 직사각형의 합”

질의가 한 번뿐이면 그냥 더하면 된다. 반복 질의가 신호다.

핵심 아이디어

처음부터 각 위치까지의 누적 합을 미리 계산해두면, 임의 구간의 합은 뺄셈 한 번이 된다.

A = [15, 13, 10,  7,  3]
S = [0, 15, 28, 38, 45, 48]   ← 한 칸 밀어 S[0]=0
        ↑           ↑
        S[1]        S[4]
 
구간 [2, 4]의 합 (2~4번째) = S[4] - S[1] = 45 - 15 = 30
                             ────────────
                             앞부분을 통째로 빼면 구간만 남는다

합 배열을 한 칸 밀어 S[0] = 0으로 두면 i = 1일 때도 예외 처리가 없다.

코드 템플릿

def build(arr):
    s = [0] * (len(arr) + 1)
    for i, x in enumerate(arr):
        s[i + 1] = s[i] + x
    return s
 
def query(s, i, j):        # 1-기반 [i, j] 구간 합
    return s[j] - s[i - 1]
 
s = build([15, 13, 10, 7, 3])
query(s, 2, 4)             # 30

itertools.accumulate를 쓰면 build가 한 줄이 된다.

from itertools import accumulate
s = [0] + list(accumulate(arr))

전처리에 O(n)이 들지만 질의 하나가 O(1)이라, 질의가 M번이면 전체가 O(n + M)이다. 매번 더하는 O(n × M)과의 차이가 곧 통과/시간 초과의 차이다.

풀이 과정 따라가기

A = [15, 13, 10, 7, 3]에서 구간 [2, 4]의 합을 구하는 과정이다.

함정

인덱스 기준 섞임. 문제는 1-기반, 코드는 0-기반인 경우가 많다. “S를 한 칸 민다”를 코드 전체에서 일관되게 지키지 않으면 off-by-one이 난다.
값이 바뀌는 문제에는 못 쓴다. 중간에 원소가 갱신되면 합 배열을 다시 만들어야 한다(O(n)). 갱신+질의가 섞이면 펜윅 트리/세그먼트 트리의 영역이다.
입력이 많다. 질의 10만 개면 input()으로도 느리다 — sys.stdin.readline을 같이 쓴다.