[Algorithm] 누적 합(prefix sum) 알고리즘 by python

누적 합(prefix sum) 알고리즘이 있습니다.

 

배열의 일부 구간에 대해 아주 빠른 속도로 구해줄 수 있는 알고리즘입니다.

 

일반적으로 배열에서 일부분의 값을 구하기 위해선 O(N)의 시간이 소요되지만, 누적 합 알고리즘을 사용하면 O(1)의 시간 복잡도를 기대해 볼 수 있습니다.

 

누적 합 알고리즘을 구하는 2가지 방법이 있습니다.

 

1. 1차원 배열일 때

 

주어진 배열 arr와 누적합을 구할 0 * (arr의 길이 + 1)로 초기화된 배열 sum_arr가 있어야 합니다.

 

그리고 주어진 배열을 순차적으로 탐색하면서 sum_arr 배열을 완성시키면 됩니다.

 

누적합 배열이기 때문에 sum_arr[i]에는 arr[0] + arr[1] + ... + arr[i-1]까지의 수의 합이 들어가 있다고 생각해야 합니다.

 

예시)

arr = [ 2, 4, 1, -5, 2, -3 ]

 

sum_arr = [ 0, 2, 6, 7, 2, 4, 1 ]

 

sum_arr의 첫 번째 인덱스는 0으로 고정입니다. sum_arr[0]을 만족하는 arr[i]가 없기 때문입니다.

 

arr의 1번째와 4번째까지의 누적합을 구할 때 sum_arr를 이용할 수 있습니다.

 

공식은 다음과 같습니다.

sum[i, j] = sum_arr[j+1] - sum_arr[i]

 

 

공식에 적용하면 sum[1, 4] = sum_arr[4+1] - sum_arr[1]가 됩니다.

 

sum_arr[5] = 4이고 sum_arr[1]은 2이므로 4 - 2 = 2가 됩니다.

 

그리고 arr = [ 2, 4, 1, -5, 2, -3 ]에서 1번째와 4번째 구간의 누적합은 2 + 4 + 1 + (-5) = 2입니다. 

 

둘의 값이 2로 동일합니다.

 

코드로 구현하면 다음과 같습니다.

for i in range(n):
    tmp += number[i]
    sum_number[i+1] = tmp

 

sum_arr[i]를 for문으로 누적합을 구해줍니다.

a, b = input().split() # 원하는 구간이 정수로 주어졌을때
print(sum_arr[b-1+1] - sum_arr[a-1]) # 인덱스가 0부터 시작하기 때문에 -1

 

2. 2차원 배열일 때

 

 

1차원 배열일 때 sum_arr의 길이가 n+1이 되었듯이 2차원 배열일때는 arr가 N X N이라면 sum_arr도 (N+1) X (M+1)이 됩니다. 

 

이때 sum_arr의 각항은 arr[0][0]부터 arr[i-1][j-1]까지의 각 항의 합이 됩니다. 

 

그림으로 알아보겠습니다.(클릭하면 출처로 이동합니다.)

 

그림에 따라 sum_arr[3][3]는 arr[0][0]부터 arr[2][2]까지의 누적합이 됩니다.

 

2차원 배열에서의 누적합을 구하는 공식은 다음과 같습니다.

sum_arr[i][j] = arr[i-1][j-1] + sum_arr[i-1][j] + sum_arr[i][j-1] - sum_arr[i-1][j-1]

 

2중 for문을 사용하여 sum_arr를 구해줄 수 있습니다.

arr = [[ 1, 2, 3, 4 ],
       [ 2, 3, 4, 5 ],
       [ 3, 4, 5, 6 ],         
       [ 4, 5, 6, 7 ]]
n = 4 # 정방형 행렬
sum_arr = [ [0] * (n+1) for _ in range(n+1) ]
for i in range(1, n+1):
    for j in range(1, n+1):
        sum_arr[i][j] = arr[i-1][j-1] + sum_arr[i-1][j] + sum_arr[i][j-1] - sum_arr[i-1][j-1]

for i in range(n+1):
    print(*sum_arr[i])

 

결과)

0 0 0 0 0
0 1 3 6 10
0 3 8 15 24
0 6 15 27 42
0 10 24 42 64

 

2차원 sum_arr를 구했으니 이제 특정 구간에서의 합을 구하는법을 알아보겠습니다.

 

arr의 (x1, y1)부터 (x2, y2)까지의 합을 sum이라고 할때

 

공식은 다음과 같습니다.

sum = sum_arr[x2+1][y2+1] - sum_arr[x1][y2+1] - sum_arr[x2+1][y1] + sum_arr[x1][y1]

 

arr(1, 1)부터 arr(3, 2)까지 더하기

sum = sum_arr[4][3] - sum_arr[1][3] - sum_arr[4][1] + sum_arr[1][1]
# 42 - 6 - 10 + 1 = 27

 

 

 

---

참고1

그림 출처

2차원 배열 누적합 문제