[Algorithm] 누적합(Prefix Sum)과 [백준 11659] 구간 합 구하기 4

6720·2023년 6월 1일

이거 모르겠어요

목록 보기

22/38

누적합

0번째 인덱스부터 N번째 인덱스까지 탐색하면서 인덱스 i일때 0번째 인덱스부터 i번째 인덱스까지의 합을 말함.

코드

import java.io.*;
import java.util.*;

public class Main {
    public static void main(String[] args) {
        int[] arr = {1, 5, 3, 7, 2, 8, 9};
        int[] prefixSum = new int[arr.length + 1];

        for (int i = 1; i < prefixSum.length; i++) {
            prefixSum[i] += arr[i-1] + prefixSum[i-1];
        }
        System.out.println(Arrays.toString(prefixSum));
    }
}

코드로 표현할 때는 원래 값이 담겨있는 배열과 누적합을 저장할 배열을 사용함.
prefixSum의 크기를 arr.length + 1로 선정하여 반복문이 조금 더 편하게 돌아가도록 설정함.
만약 prefixSum에서 값을 불러와야 한다면 본래 인덱스 + 1을 인덱스로 사용하면 됨.

[결과]

[0, 1, 6, 9, 16, 18, 26, 35]

누적합과 [백준 11659] 구간 합 구하기 4

다음의 문제는 겉보기에는 반복문으로 첫 인덱스와 끝 인덱스를 설정하여 간단하게 답을 구할 수 있는 문제임. 실제로도 그렇지만 다음과 같은 조건이 있음.

“시간 제한”

시간 제한이라는 문구만 보면 자동으로 튀어나오는 말, “반복문은 자제하라”

필자의 경우는 시간 제한 문제가 나오면 자동으로 다이나믹 프로그래밍(dp)가 떠오르도록 학습되어 있는 상태였음. 하지만 dp와 해당 문제는 연관성이 전혀 없어보였음.

그렇게 검색하다가 알아낸 알고리즘이 누적합이었음.

문제 접근

prefixSum 배열

해당 문제는 누적합을 저장하는 prefixSum 배열을 생성하면 그 후에는 원래 배열 arr의 쓸모가 없어짐.
→ 애초에 prefixSum으로 저장되도록 하면 어떨까?

[prefixSum 배열 하나로만 푸는 코드]

public class Main {
    static int[] prefixSum;
    public static void main(String[] args) throws IOException {
        ...
        StringTokenizer st = new StringTokenizer(br.readLine());
        int n = Integer.parseInt(st.nextToken());
        int m = Integer.parseInt(st.nextToken());
        prefixSum = new int[n+1];

        st = new StringTokenizer(br.readLine());
        for (int i = 1; i <= n; i++) {
            prefixSum[i] = Integer.parseInt(st.nextToken()) + prefixSum[i-1];
        }
}

입력값을 직접 받아와야 한다는 백준의 특징을 착안해서 다음처럼 작성하였음.
prefixSum 배열이 반복문 안에서만 해결되도록 prefixSum의 크기를 n+1로 설정하였음.
인덱스가 0인 곳의 값이 0으로 나오기도 하고 인덱스가 한 칸씩 밀리기도 하지만, 하나의 반복문에서 해결된다는 장점도 만만치 않음.

다만 Integer.parseInt(st.nextToken())을 사용해 연산을 하면 시간을 더 잡아먹기는 하는 듯.
이 부분은 되도록 int num 등 하나의 변수로 뺀 다음에 연산하기도 하지만 오늘따라 한 줄로 끝나는 코드를 작성하고 싶었음.

누적합을 통해 특정 범위의 합만 구하기

해당 문제에서의 입력값 인덱스에는 +1이 되어 있음.
예를 들어서 arr에서는 0~4 인덱스라면 입력 인덱스는 1~5가 됨.
현재는 prefixSum의 크기가 arr의 크기 + 1이라서 따로 신경써줄 필요는 없음.
어떻게 생각하면 prefixSum의 크기를 +1 한건 좋은 선택이었을지도 모름.

반환하는 값의 조건에는 시작 인덱스에 따라 값이 달라짐.
startIdx가 1인, 즉 endIdx 인덱스까지 모든 값을 더해야 할 경우 prefixSum[end]을 바로 반환해주면 됨.
하지만 startIdx가 1이 아닌 경우, 누적합에서 앞 부분에 포함하지 않는 숫자가 생기면 그 부분에 대한 누적합을 빼줘야 함.
이때 빼줘야 하는 값의 인덱스는 startIdx - 1 부분이 됨.

이렇게 조건에 대한 코드를 작성했을 때 하나 생각나는 것이 있을 것임.
startIdx가 1일 때도 return prefixSum[endIdx] - prefixSum[startIdx - 1]을 사용해도 되지 않을까?
왜냐하면 startIdx가 1일 때는 prefixSum[startIdx - 1]은 prefixSum[0]을 가리키고, 해당 값은 0이기 때문임.
그러므로 어떤 조건에도 상관없이 prefixSum[endIdx] - prefixSum[startIdx - 1]를 반환해주면 됨.

조금 더 나아가서 startIdx를 배열에 입력되기 전에 -1해서 넣어줄 수도 있음.
이렇게 하면 조금 보기 편해질 것임.

import java.io.*;
import java.util.*;

public class Main {
    static int[] prefixSum;
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        StringTokenizer st = new StringTokenizer(br.readLine());
        int n = Integer.parseInt(st.nextToken());
        int m = Integer.parseInt(st.nextToken());
        prefixSum = new int[n+1];

        st = new StringTokenizer(br.readLine());
        for (int i = 1; i <= n; i++) {
            prefixSum[i] = Integer.parseInt(st.nextToken()) + prefixSum[i-1];
        }

        while (m-- > 0) {
            st = new StringTokenizer(br.readLine());
            int start = Integer.parseInt(st.nextToken()) - 1;
            int end = Integer.parseInt(st.nextToken());
            bw.write(prefixSum[end] - prefixSum[start] + "\n");
        }
        bw.flush();
        br.close();
        bw.close();
    }
}