[Java Algorithm] 누적 합(prefix sum)과 2차원 누적 합

누적 합은 누적 합이 뭔지 설명하는 것보다 코드를 보면서 이해하는 것이 더 쉽기 때문에 문제를 예로 들어보자

---

📝 문제

크키가 N인 정수 배열 A가 있고, 여기서 다음과 같은 연산을 최대 M번 수행해야 하는 문제가 있다.

• 구간 l, r ( l <= r ) 이 주어져 있을 때, A[l] + A[l + 1] + ... + A[r -1] + A[r]을 구해서 출력하기

이때 위의 합을 구하기 위해 반복문을 사용하여 푸는 방법이 있다.

 

📍 반복문을 이용하여 문제 해결

int answer = 0;

for(int i = l; i <= r; i++) {
	answer += a[i];
}

위 코드의 시간 복잡도는 O(N) 이다. 문제에서는 최대 M번을 진행해야 하기 때문에 총 시간 복잡도는 O(NM)이 된다.

그렇다

 

❗️누적 합

누적 합 아이디어는 배열 A에 들어있는 값이 바뀌지 않는다는 점을 이용한다.

배열이 변하지 않으니 구간의 합도 변하지 않는다. 따라서 앞에서부터 차례대로 누적된 합을 구해놓고 이를 이용해서 구간의 합을 구할 수 있다.

 

S[i] = A[1]+ ... + A[i], S[0] = 0 으로 정의해보자

 

l 번째 수부터 r 번째 수까지 합은 S[r] - S[l - 1]과 같다.

이유를 살펴 보자 S[r]  = A[1] + ... + A[r] 이고 S[l - 1] = A[1] + ... + A[l - 1] 이다.

 

여기서 r이 더 크기 때문에 S[r] - S[l - 1]은 A[l] + ... + A[r] 이 된다.

 

이 때문에 구간의 합을 구하기 위해서 뺄셈 연산 한 번만 하면 되기 때문에 시간 복잡도는 O(1)이며 총 M번 수행하기 때문에 O(M)이 된다.

 

i	0	1	2	3	4	5	6	7	8
A[i]		7	2	4	5	9	1	9	2
S[i]	0	7	9	13	18	27	28	37	39
S[r]	0	7	9	13	18	27	28	37	39
S[l - 1]	0	7	9	13	18	27	28	37	39
S[r] - S[l - 1]	0	7	9	13	18	27	28	37	39

위 예시는 l = 3, r = 7인 경우를 예로 들었다.

 

 

---

✏️ 구현

수 N개가 주어졌을 때, i번째 수부터 j번째 수까지 합을 구하는 프로그램을 작성하시오.

라는 문제가 주어지고 입력 주건으로 첫째 줄에 수의 개수 N과 합을 구해야 하는 횟수 M이 주어진다. 둘째 줄에는 N개의 수가 주어진다. 수는 1,000보다 작거나 같은 자연수이다. 셋째 줄부터 M개의 줄에는 합을 구해야 하는 구간 i와 j가 주어졌을 때 아래와 같이 문제를 해결할 수 있다.

 

import java.util.*;
public class Main {
    public static void main(String args[]) {
        Scanner sc = new Scanner(System.in);
        int n = sc.nextInt();
        int m = sc.nextInt();
        int[] a = new int[n+1];
        int[] s = new int[n+1];
        for (int i=1; i<=n; i++) {
            a[i] = sc.nextInt();
            s[i] = s[i-1] + a[i];
        }
        while (m-- > 0) {
            int x = sc.nextInt();
            int y = sc.nextInt();
            System.out.println(s[y]-s[x-1]);
        }
    }
}

---

❗️추가 정리

2차원으로 생각하기 전 1차원으로 먼저 생각해보자

기본 배열으로 [1, 2, 3, 4, 5]이 있다고 가정한 후 2 ~ 3 까지는 2를 빼주고 0 ~ 2까지는 4를 빼준다 라는 문제가 있다.

 

이 때 가장 먼저 누적합을 저장할 새로운 배열 [0, 0, 0, 0, 0]을 만들어 준다.

 

이제 2 ~ 3까지 2를 빼준다는 [0, 0, -2, 0, 2] 이렇게 표현할 수 있다. 

여기서 l 이 2가 되고 r이 3이 되며 A[l] 은 변화 값인 -2를, 누적합이 변하면 안되기 때문에 A[r + 1] 에는 +2를 하여 누적합에 변함이 없게 만든다.

 

두번째로 0 ~ 2까지는 4를 빼준다에서 l 이 0이 되고 r이 2가 되기 때문에 A[0] = -4, A[2 + 1] = A[3] = 4가 된다.

 

따라서 정리해보자면 [-4, 0, -2, 4, 2] 라는 배열이 만들어 지고 누적 합을 사용하게 되면

 

[-4, -4, -6, -2, 0] 이라는 차들의 합으로 나타낼 수 있다.

 

여기서 해당 배열을 [1, 2, 3, 4, 5]와 합쳐준다면 문제를 좀더 빠르게 해결할 수 있다.