어떤 정수들의 배열이 주어졌을 때 그 배열의 연속된 일부의 총 합을 매우 빠르게 구할 수 있는 방법을 배웠다. 바로 Prefix Sums 를 사용하는 방법이다.
Prefix Sums
Prefix Sums 란 어떤 정수 배열 내 요소들의 각 index 까지의 총 합을 저장해놓은 별도의 배열을 말한다.
예를 들어서 아래와 같은 배열 a가 있을 때,
a = [1, 2, 3, 4, 5, 6]
이 배열의 Prefix Sums 배열 P는 다음과 같다.
P = [0, 1, 3, 6, 10, 15, 21]
이는 아래와 같은 계산을 통해 구해진다.
P[0] = 0
P[1] = a[0]
P[2] = a[0] + a[1]
P[3] = a[0] + a[1] + a[2]
P[4] = a[0] + a[1] + a[2] + a[3]
.
.
.
이걸 다시 쓰면 다음과 같이 된다.
P[0] = 0
P[1] = P[0] + a[0]
P[2] = P[1] + a[1]
P[3] = P[2] + a[2]
.
.
.
P[k] = P[k - 1] + a[k - 1]
Prefix Sums 의 계산을 함수로 작성하면 다음과 같다.
def prefix_sums(A):
n = len(A)
P = [0] * (n + 1)
for k in range(1, n + 1):
P[k] = P[k - 1] + A[k - 1]
return P
Prefix Sums 활용하기
이 Prefix Sums 를 활용하면 다음과 같은 배열의 슬라이스들의 총 합을 매우 빠르게 구할 수 있다.
[1, 2, 3]
[2, 3, 4]
[3, 4, 5, 6]
.
.
.
각각의 슬라이스들의 총 합을 구하는 가장 심플한 방법은 각 슬라이스들을 순회하면서 각 요소들을 더해주면 된다.
sum([1, 2, 3])
sum([2, 3, 4])
sum([3, 4, 5, 6])
.
.
.
만약에 위와 같은 슬라이스 m 개의 각각의 총 합을 구해야한다고 생각해보자. 이때 가장 간단한 접근방법으로 계산한다면 최악의 경우 \( O(nm) \) 시간이 걸린다.
하지만 Prefix Sums 를 활용하면 \( O(n+m) \) 시간만에 해낼 수 있다.
a[0:2] 인 슬라이스 [1, 2, 3] 의 총 합을 Prefix Sums 를 이용해서 구해보자.
Prefix Sums P는 아래와 같다.
P = [0, 1, 3, 6, 10, 15, 21]
a[0]에서 a[2] 까지의 총 합은 P[3] 이므로 P[3] 은 sum([1, 2, 3]) 과 같다. 하지만 실행시간은 \( O(1) \) 으로 훨씬 빠르다.
a[2:5] 인 슬라이스 [3, 4, 5] 의 총 합은 어떻게 구할까?
생각을 해보면 [3, 4, 5] 의 총 합은 [1, 2, 3, 4, 5] 의 총 합에서 [1, 2] 의 총 합을 뺀 값이다. 각각의 값을 P에서 찾아서 계산해주면 된다.
[1, 2, 3, 4, 5] 의 총 합은 P[5], [1, 2] 의 총 합은 P[2] 이다. 따라서 P[5] - P[2] 를 해주면 [3, 4, 5] 의 총 합을 구할 수 있다.
일반화하면 a[x:y] 의 총 합은 P[y] - P[x] 와 같다.
이것을 함수로 정의하면 아래와 같다.
def count_total(P, x, y):
return P[y] - P[x]
이를 활용해 m 개의 슬라이스들의 합을 구하면 \( O(n+m) \) 시간만에 해낼 수 있다.