[Algorithm] 누적합과 구간합

누적합

• 배열의 앞에서부터 특정 인덱스까지의 합을 미리 계산해두는 기법
• 빠르게 구간합을 구할 수 있음
• 일반적으로 prefix[i] = prefix[i-1] + arr[i] 형태로 계산

구간합

• 특정 구간 [L, R] 의 합을 구하는 문제
• 누적합을 사용하면 prefix[R] - prefix[L -1] 로 빠르게 계산 가능

예제

arr = [1, 2, 3, 4, 5]
n = len(arr)

prefix_sum = [0] * (n + 1)
for i in range(1, n + 1):
    prefix_sum[i] = prefix_sum[i - 1] + arr[i - 1]
    
def range_sum(left, right):
    return prefix_sum[right] - prefix_sum[left - 1]

print(range_sum(2, 4))

2차원 누적합

개념

• 2차원 배열에서 사각형 범위의 합을 빠르게 구하는 기법
• 1차원 누적합처럼 미리 합을 계산해두고, 구간의 합을 빠르게 찾음
• 특정 구간 [x1, y1] ~ [x2, y2] 의 합을 O(1) 에 구할 수 있음

누적합 배열 생성

현재 위치까지의 누적합은 위쪽 + 왼쪽 - 중복된 부분 + 현재 값으로 계산

구간합 계산

위에서 빼고, 왼쪽에서 빼고, 중복된 부분을 다시 더함

예제

# 2D 배열 (5x5 예제)
arr = [
    [3, 7, 9, 2, 1],
    [5, 8, 4, 3, 6],
    [9, 2, 3, 7, 8],
    [4, 6, 5, 8, 2],
    [1, 3, 7, 9, 4]
]

# 행과 열 크기
n, m = len(arr), len(arr[0])

# 2차원 누적합 배열
prefix = [[0] * (m + 1) for _ in range(n + 1)]

# 누적합 계산
for i in range(1, n + 1):
    for j in range(1, m + 1):
        prefix[i][j] = (
            arr[i - 1][j - 1] 
            + prefix[i - 1][j] 
            + prefix[i][j - 1] 
            - prefix[i - 1][j - 1]
        )

# 특정 구간 (x1, y1) ~ (x2, y2)의 합 구하기
def range_sum(x1, y1, x2, y2):
    return (
        prefix[x2][y2] 
        - prefix[x1 - 1][y2] 
        - prefix[x2][y1 - 1] 
        + prefix[x1 - 1][y1 - 1]
    )

# (2,2) ~ (4,4) 구간 합
print(range_sum(2, 2, 4, 4))