Post

구간 합 (prefix sum)

목차


개요

구간 합 알고리즘은 n-m구간 사이의 합을 빠르게 구하는 알고리즘이다.

내용 및 백준 출처 : https://www.crocus.co.kr/843


문제

예를 들어

1
data = [ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]

배열이 존재한다고 가정하자. 총 길이 12의 배열에서 n 과 m 사이의 구간 합을 요구하는 쿼리가 1천만 개가 들어왔을 때 가장 쉬운 구현 방식은 다음과 같다.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
data = [ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
query_count = 10000000      #1천만개

# random module import
from random import randint
N_MAX_RANDOM_VAL = len(data) - 1
M_MAX_RANDOM_VAL = len(data)

for q in range(query_count):
    n = randint(0, N_MAX_RANDOM_VAL)
    m = randint(n + 1, M_MAX_RANDOM_VAL)

    _sum = 0
    for i in range(n, m):
        _sum += data[i]
    print("{} ~ {} sum : {}".format(n, m, _sum))

해당 코드의 시간 복잡도는 query_count (10,000,000) * (최악의 시간을 가정 : n = 0, m = 1212번의 반복문) = 120,000,000
= O(query_count * m)

query_count 1 기준
= O(m-n)

해당 시간복잡도를 보듯 nm 사이의 범위가 넓어질수록 시간 복잡도가 증가하게 될 것이다.
data 배열에 1억개의 데이터가 들어있는 때를 가정하면 최악의 경우 1쿼리당 1억번 반복문을 돌아야 할 것이다.


해결방법

0~i 까지의 합을 저장하는 배열을 만든다.

1
prefix_sum = [ data[0], data[0~1], data[0~2], data[0~3], data[0~4] ... ]

이후 n~m 구간의 합을 구할 때 prefix_sum[m] - prefix_sum[n - 1] 을 하면
0~m - 0~(n-1) 이 되므로 n~m의 합을 구할 수 있다.

그림으로 표현한 구간 합

1
2
3
4
5
6
7
8
9
n = 4, m = 6

arr index   :|0|1|2|3|4|5|6|
             |-|-|-|-|-|-|-|
prefix[m]   =|1|2|3|4|5|6|7|= 28
prefix[n-1] =|1|2|3|4| | | |= 10

result
             | | | | |5|6|7|= 18

구간 합 코드

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
data = [ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]

# 배열 합 저장하는 배열 생성
prefix_sum = [ data[0] ]

for i in range(1, len(data)):
    # 이전까지의 합에 data[i] 를 더해 0~i 까지의 합을 저장
    sum_value.append(sum_value[i - 1] + data[i])

query_count = 10000000      #1천만개

# random module import
from random import randint
N_MAX_RANDOM_VAL = len(data) - 1
M_MAX_RANDOM_VAL = len(data)

for q in range(query_count):
    n = randint(0, N_MAX_RANDOM_VAL)
    m = randint(n + 1, M_MAX_RANDOM_VAL)

    print("{} ~ {} sum : {}".format(n, m, prefix_sum[m] - prefix_sum[n - 1]))

이 경우 prefix_sum[m] - prefix_sum[n - 1] 빼기 연산 한 번만 진행하면 결과를 알 수 있으므로
시간 복잡도는 O(query_count) 이다.

query_count 1 기준
= O(1)

처음 사용하였던 로직과 달리 nm 사이의 범위가 아무리 넓어도 시간 복잡도는 O(1)이다.


Prefix Sum이 쓰이는 문제들

[11441번] 합 구하기 : https://www.acmicpc.net/problem/11441

[1806번] 부분합 : http://www.crocus.co.kr/595

[1644번] 소수의 연속합 : http://www.crocus.co.kr/599

[2015번] 수들의 합 4 : http://www.crocus.co.kr/602

[Codeground 42번] 부분배열 : http://www.crocus.co.kr/841

This post is licensed under CC BY 4.0 by the author.