[수학] 누적 합 (Python)

길이가 N인 수열 $\left\{ a_n \right\}$에서 아래의 연산을 $M$번 수행하시오.
• 두 정수 $i$, $j$($i \leq j$)를 입력 받아 $\sum_{k=i}^{j}a_k$ ($=a_i + a_{i+1} + ... + a_j$)를 구하여라.

 

방법1: 매번 다 더해보기

문제에서 요구하는 연산은 $a_i$부터 $a_j$까지 다 더하면 되는 것으로, 구현하기 어렵지 않다.

바로 아래와 같이 코드를 짜볼 수 있겠다.

for _ in range(M):
    i, j = map(int, input().split())
    answer = 0
    
    for k in range(i, j + 1):
        answer += a[k - 1]
    
    print(answer)

 

첫째 항인 $a_1$이 리스트 a에서는 a[0]이므로. $a_k$를 더할 때 a[k - 1]을 더해줘야 한다.

최악의 경우는 매번 첫번째 항부터 마지막 항까지 더하는 것으로, 시간 복잡도는 $O(NM)$이 된다.

 

방법2: 미리 더해놓기 (누적 합)

수열 $\left\{a_n\right\}$을 토대로 새로운 수열 $\left\{S_n\right\}$을 만들어보자.

$S_n = a_1 + a_2 + ... + a_n = \sum_{k=1}^{n}a_k$이다. 즉, $S_n$은 $a_1$부터 $a_n$까지의 합을 의미한다.

 

이제 우리가 구하려는 값을 $\left\{S_n\right\}$을 이용한 식으로 아래와 같이 변형시킬 수 있다.

$\sum_{k=i}^{j}a_k = \sum_{k=1}^{j} - \sum_{k=1}^{i-1} = S_j - S_i (i \geq 2)$

($i=1$인 경우, $\sum_{k=i}^{j}a_k = S_j$이다.)

 

이제 우리가 구하고자 하는 식에서 $\sum$이 사라졌다.

이제 매번 반복문을 돌지 않고 뺄셈 연산 한 번으로 값을 구할 수 있다.

(물론 $\left\{S_n\right\}$을 만들 때 한 번은 반복문을 돌아야 한다.)

 

구현해보기

S[0] = a[0]

for k in range(1, N):
    S[k] = S[k - 1] + a[k]

for _ in range(M)
    i, j = map(int, input().split())
    
    if i == 1:
        print(S[j - 1])
    else:
        print(S[j - 1] - S[i - 2])

 

위 이론대로 리스트 S를 만들어 구현하였다. (선언하는 부분은 생략하였다.)

이제 여기에서 몇 가지 부분을 수정해 최종 코드를 짜보겠다.

 

1. 수열에서의 번호와 리스트의 인덱스가 불일치하는 문제를 해결하기 위해, 리스트 a의 맨 앞에 0을 추가해준다.
   이렇게 하면 i가 1로 들어오는 경우도 따로 처리할 필요가 없어진다.
2. 굳이 새로운 리스트 S를 만들지 않고, 리스트 a를 재활용해도 실행에는 지장이 없다.

 

따라서 입출력을 포함한 최종 코드는 아래와 같다.

 

최종 코드

import sys
input = sys.stdin.readline

N, M = map(int, input().split())
a = [0] + list(map(int, input().split()))

for k in range(1, N + 1):
    a[k] += a[k - 1]

for _ in range(M):
    i, j = map(int, input().split())
    print(a[j] - a[i - 1])

 

리스트 a의 맨 앞에 0을 추가함에 따라, 누적합을 계산할 때 a[N-1]이 아닌 a[N]까지 계산해야 함에 주의하자.

이렇게 구현한 누적 합 알고리즘의 시간 복잡도는 $O(N+M)$이 된다.

 

이차원에서의 누적 합

지금까지 수열(일차원 배열)에서의 구간 합을 구하는 방법 중 하나인 누적 합에 대해 알아보았다.

그렇다면 행렬(이차원 배열)에서의 구간 합을 구할 때에도 누적 합을 사용할 수 있을까?

 

$a_{i_1i_2}$부터 $a_{j_1j_2}$까지의 구간 합은 $\sum_{k_1=i_1}^{j_1} \sum_{k_2=i_2}^{j_2} a_{k_1k_2}$로 나타낼 수 있으며, 그 범위는 아래와 같다.

 

$a_{11}$	...	$a_{1i_2}$	...	$a_{1j_2}$
⋮		⋮		⋮
$a_{i_11}$	...	$a_{i_1i_2}$	...	$a_{i_1j_2}$
⋮		⋮		⋮
$a_{j_11}$	...	$a_{j_1i_2}$	...	$a_{j_1j_2}$

 

이제 일차원과 마찬가지로 새로운 행렬 S를 만들어보자.

$S_{n_1n_2} = \sum_{k_1=1}^{n_1} \sum_{k_2=1}^{n_2} a_{k_1k_2}$로 정의하면 된다.

이제 $S_{(i_1-1)(i_2-1)}$와 $S_{j_1j_2}$만 있으면 원하는 값을 구할 수 있는지 확인해보자. 두 구간을 표시하면 아래와 같다. (각각 빨간색, 하늘색이다.)

 

$a_{11}$	...	$a_{1i_2}$	...	$a_{1j_2}$
⋮		⋮		⋮
$a_{i_11}$	...	$a_{i_1i_2}$	...	$a_{i_1j_2}$
⋮		⋮		⋮
$a_{j_11}$	...	$a_{j_1i_2}$	...	$a_{j_1j_2}$

 

두 영역을 보면 이 영역만으로는 구하고자 하는 값을 얻어내기 힘들다.

여기에 두 개의 값만 더 추가하면 원하는 영역의 값을 구할 수 있다. 바로 $S_{j_1(i_2-1)}$와 $S_{(i_1-1)j_2}$이다.

두 구간을 각각 주황색과 초록색으로 칠하면 아래와 같다. (겹치는 부분은 노란색으로 칠했다.)

 

$a_{11}$	...	$a_{1i_2}$	...	$a_{1j_2}$
⋮		⋮		⋮
$a_{i_11}$	...	$a_{i_1i_2}$	...	$a_{i_1j_2}$
⋮		⋮		⋮
$a_{j_11}$	...	$a_{j_1i_2}$	...	$a_{j_1j_2}$

 

위 표에서 한 번도 포함되지 않은 영역이 바로 우리가 구하고자 하는 영역이다.

즉, 전체 영역(하늘색)에서 위 두 영역을 빼면 된다. 식으로 나타내면 아래와 같다.

$S_{j_1j_2} - S_{j_1(i_2-1)} - S_{(i_1-1)j_2}$

 

아직 식이 완성되지 않았다. 아까 세 번째 표에서 겹쳤던 부분(노란색)이 두 번 빼졌기 때문이다.

해당 영역은 두 번째 표에서 빨간색 영역과 일치하므로, 위의 식에서 해당 부분을 한 번 더해주면 된다.

따라서 완성된 식은 아래와 같다.

$\sum_{k_1=i_1}^{j_1} \sum_{k_2=i_2}^{j_2} a_{k_1k_2} = S_{j_1j_2} - S_{j_1(i_2-1)} - S_{(i_1-1)j_2} + S_{(i_1-1)(i_2-1)}$

 

S를 구하는 방법

일차원에서의 누적합을 구하듯이 가로랑 세로 방향 모두를 고려해 누적합을 구해주면 된다.

이때, 대각성 방향의 경우 두 번 더해지므로 한 번 빼줘야 한다. 아래는 예시 코드이다.

for k1 in range(1, N1):
    for k2 in range(1, N2):
        S[k1][k2] = a[k1][k2] + S[k1 - 1][k2] + S[k1][k2 - 1] - S[k1 - 1][k2 - 1]

 

일차원과 마찬가지로 원본 배열이 필요 없다면 S를 굳이 만들지 않아도 된다.

for k1 in range(1, N1):
    for k2 in range(1, N2):
        a[k1][k2] += a[k1 - 1][k2] + a[k1][k2 - 1] - a[k1 - 1][k2 - 1]

 

최종 코드

import sys
input = sys.stdin.readline

N1, N2, M = map(int, input().split())
a = [[0] * (N2 + 1)]

for _ in range(N1):
    a.append([0] + list(map(int, input().split())))

for k1 in range(1, N1 + 1):
    for k2 in range(1, N2 + 1):
        a[k1][k2] += a[k1 - 1][k2] + a[k1][k2 - 1] - a[k1 - 1][k2 - 1]

for _ in range(M):
    i1, i2, j1, j2 = map(int, input().split())
    print(a[j1][j2] - a[i1 - 1][j2] - a[j1][i2 - 1] + a[i1 - 1][i2 - 1])

 

관련 문제 (BOJ)

• [11659] 구간 합 구하기 4: 누적 합 기법을 그대로 사용하면 되는 문제
• [11660] 구간 합 구하기 5: 2차원에서의 누적 합 기법을 그대로 사용하면 되는 문제
• [10986] 나머지 합: 누적 합을 활용하는 응용 문제. N의 최댓값이 크므로 수학적인 생각이 필요하다.
• [25682] 체스판 다시 칠하기 2: 2차원에서의 누적 합 기법을 활용하는 응용 문제. 문제를 어떻게 숫자로 바꿀 수 있을까?