파이썬 2차원 배열에서의 누적합, 구간합

2차원 배열에서의 누적합(Prefix Sum)을 계산하고 이를 활용해 구간합을 효율적으로 구하는 방법은 다음과 같습니다.

1. 2차원 배열의 누적합 계산

2차원 배열에서의 누적합은 prefix[i][j]를 다음과 같이 정의합니다:

prefix[i][j]=원점 (0,0)부터 (i,j)까지의 모든 원소의 합

계산식

prefix[i][j] = array[i][j] + prefix[i−1][j] + prefix[i][j−1] − prefix[i−1][j−1]

• array[i][j]: 현재 값
• prefix[i−1][j]: 위쪽 영역의 누적합
• prefix[i][j−1]: 왼쪽 영역의 누적합
• prefix[i−1][j−1]: 중복된 영역(대각선 방향)의 누적합을 빼줌

2. 2차원 구간합 계산

구간합을 구하고자 하는 범위가 (x1,y1)부터 (x2,y2)라고 할 때:

계산식

sum[x1:x2, y1:y2] = prefix[x2][y2] − prefix[x1−1][y2] − prefix[x2][y1−1] + prefix[x1−1][y1−1]

• prefix[x2][y2]: 전체 영역의 누적합
• prefix[x1−1][y2]: 위쪽 제외
• prefix[x2][y1−1]: 왼쪽 제외
• prefix[x1−1][y1−1]: 중복 제외

3. 알고리즘 구현 (Python)

def compute_prefix_sum(matrix):
    rows, cols = len(matrix), len(matrix[0])
    prefix = [[0] * cols for _ in range(rows)]
    
    for i in range(rows):
        for j in range(cols):
            prefix[i][j] = matrix[i][j]
            if i > 0:
                prefix[i][j] += prefix[i-1][j]
            if j > 0:
                prefix[i][j] += prefix[i][j-1]
            if i > 0 and j > 0:
                prefix[i][j] -= prefix[i-1][j-1]
    return prefix

def range_sum(prefix, x1, y1, x2, y2):
    result = prefix[x2][y2]
    if x1 > 0:
        result -= prefix[x1-1][y2]
    if y1 > 0:
        result -= prefix[x2][y1-1]
    if x1 > 0 and y1 > 0:
        result += prefix[x1-1][y1-1]
    return result

# 예제
array = [
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9]
]
prefix = compute_prefix_sum(array)
print(range_sum(prefix, 1, 1, 2, 2))  # 결과: 28