Python/자료구조

누적합 배열 (Prefix Sum Array) 자료구조

metamong-data 2025. 3. 10. 19:17
728x90
반응형

📌 개요

누적합(Prefix Sum) 배열은 배열의 특정 구간 합을 빠르게 계산하기 위한 자료구조입니다.
일반적으로 배열에서 여러 개의 구간 합을 반복적으로 구해야 하는 문제에서 유용합니다.

누적합 배열을 사용하는 이유:

  • 단순한 구간 합 계산은 O(N) → 여러 번 반복하면 비효율적
  • 누적합 배열을 이용하면 O(1)로 구간 합을 계산할 수 있음
  • 미리 계산한 값을 활용하여 빠른 질의(Query) 처리 가능

📌 누적합 배열 개념

배열 A누적합 배열 S 를 다음과 같이 정의합니다.

S[i]=A[0]+A[1]+...+A[i]S[i] = A[0] + A[1] + ... + A[i]

즉, S[i]는 배열 A의 0번 인덱스부터 i번 인덱스까지의 합을 의미합니다.

✅ 특정 구간 [L, R] 의 합을 구하는 공식
sum(L,R)=S[R]−S[L−1]\text{sum}(L, R) = S[R] - S[L-1]
(단, L == 0이면 S[L-1] 대신 0 사용)


📌 예제 코드 (누적합 배열 생성 & 구간 합 구하기)

1. 누적합 배열 생성

def prefix_sum(arr):
    n = len(arr)
    S = [0] * (n + 1)  # 인덱스를 1부터 사용하기 위해 크기 n+1
    for i in range(1, n + 1):
        S[i] = S[i - 1] + arr[i - 1]  # 누적합 계산
    return S

arr = [3, 2, 4, 1, 7, 5]
S = prefix_sum(arr)
print(S)  # 출력: [0, 3, 5, 9, 10, 17, 22]

시간 복잡도: O(N)
배열 S를 만들면 특정 구간의 합을 O(1)에 구할 수 있음


2. 구간 합 빠르게 구하기

def range_sum(S, L, R):
    return S[R] - S[L - 1]  # 구간 합 공식 사용

# 예제
arr = [3, 2, 4, 1, 7, 5]
S = prefix_sum(arr)
print(range_sum(S, 2, 5))  # 2번부터 5번 인덱스까지의 합 (출력: 14)

시간 복잡도: O(1)
여러 개의 구간 합을 빠르게 구할 수 있음


📌 응용 문제

1. 부분합이 특정 값이 되는 구간 개수 찾기

💡 문제:
배열 arr에서 부분합이 K가 되는 구간의 개수를 찾는 문제

🔹 해결 방법:

  • 누적합 배열을 사용하여 O(N) 시간 복잡도로 해결 가능
  • 해시맵을 활용하여 이전에 등장한 누적합을 저장
from collections import defaultdict

def count_subarrays_with_sum(arr, K):
    S = 0
    count = 0
    prefix_counts = defaultdict(int)
    prefix_counts[0] = 1  # S[L-1] = 0인 경우 고려

    for num in arr:
        S += num  # 누적합 계산
        count += prefix_counts[S - K]  # S[L-1] == S[R] - K인 경우 찾기
        prefix_counts[S] += 1  # 현재 누적합 저장

    return count

arr = [1, 2, 3, 4, 5]
K = 5
print(count_subarrays_with_sum(arr, K))  # 출력: 2 ([2,3], [5])

시간 복잡도: O(N)
해시맵을 활용하여 특정 누적합이 존재하는지 빠르게 찾을 수 있음


2. 2차원 누적합 (Prefix Sum 2D)

💡 문제:
2차원 배열에서 특정 영역 [x1, y1] ~ [x2, y2] 의 합을 빠르게 구하는 문제

🔹 해결 방법:

  • 2차원 누적합 배열을 활용하면 O(1)에 구간 합을 계산 가능
  • 점화식: S[x][y]=A[x][y]+S[x−1][y]+S[x][y−1]−S[x−1][y−1]S[x][y] = A[x][y] + S[x-1][y] + S[x][y-1] - S[x-1][y-1]
    sum(x1,y1,x2,y2)=S[x2][y2]−S[x1−1][y2]−S[x2][y1−1]+S[x1−1][y1−1]\text{sum}(x1, y1, x2, y2) = S[x2][y2] - S[x1-1][y2] - S[x2][y1-1] + S[x1-1][y1-1]
def prefix_sum_2d(matrix):
    rows, cols = len(matrix), len(matrix[0])
    S = [[0] * (cols + 1) for _ in range(rows + 1)]

    for i in range(1, rows + 1):
        for j in range(1, cols + 1):
            S[i][j] = matrix[i-1][j-1] + S[i-1][j] + S[i][j-1] - S[i-1][j-1]

    return S

def range_sum_2d(S, x1, y1, x2, y2):
    return S[x2][y2] - S[x1-1][y2] - S[x2][y1-1] + S[x1-1][y1-1]

# 예제
matrix = [
    [3, 7, 1, 4],
    [2, 5, 6, 8],
    [1, 9, 2, 3]
]
S = prefix_sum_2d(matrix)
print(range_sum_2d(S, 1, 1, 2, 3))  # 특정 부분의 합 출력

시간 복잡도:

  • 누적합 배열 생성: O(N × M)
  • 질의(Query) 처리: O(1)

📌 누적합 자료구조의 장점과 단점

장점

  • 여러 개의 구간 합을 O(1)로 빠르게 계산 가능
  • 2차원 배열에도 확장 가능하여 다양한 문제 해결 가능
  • 전처리 과정이 O(N) 또는 O(N × M)으로 빠름

단점

  • 배열 크기가 커지면 메모리 사용량 증가 (특히 2차원 배열)
  • 실시간 업데이트가 필요한 경우 느림 → 세그먼트 트리(Segment Tree) 필요

📌 결론

  • 누적합(Prefix Sum) 배열을 활용하면 특정 구간 합을 빠르게 구할 수 있음
  • 1차원뿐만 아니라 2차원에서도 활용 가능
  • 여러 개의 질의를 처리해야 할 때 매우 유용한 기법
    🚀 누적합은 "구간 합을 자주 계산해야 하는 문제"에서 필수적으로 알아야 하는 자료구조입니다!
728x90