[Java] 백준 11659번 문제(구간 합 구하기 4)

문제설명

 

소스코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.StringTokenizer;

public class Boj_11659 {
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        String NM = br.readLine();
        StringTokenizer st = new StringTokenizer(NM);
        StringBuilder sb = new StringBuilder();
        int N = Integer.parseInt(st.nextToken());
        int M = Integer.parseInt(st.nextToken());
        String inputArr = br.readLine();
        st = new StringTokenizer(inputArr);

        int[] numArr = new int[N];
        for (int i = 0; i < numArr.length; i++) {
                numArr[i] = Integer.parseInt(st.nextToken());
        }

        int[] prefixSum = new int[N + 1];
        for (int i = 1; i <= N; i++) {
            prefixSum[i] = numArr[i - 1] + prefixSum[i-1];
        }
        for(int i = 0; i < M ; ++i){
            String inputTwo = br.readLine();
            st = new StringTokenizer(inputTwo);
            int start = Integer.parseInt(st.nextToken());
            int end = Integer.parseInt(st.nextToken());
            int sum = prefixSum[end] - prefixSum[start - 1];

            sb.append(sum).append("\n");
        }
        System.out.println(sb);
    }
}

 

설명

이 문제의 시간 초과를 방지하는 핵심 아이디어는 누적 합 배열을 구하는 것이다.

 

먼저, BufferedReader를 사용하여 입력을 받고, 이를 처리한다.

BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
String NM = br.readLine();
StringTokenizer st = new StringTokenizer(NM);
StringBuilder sb = new StringBuilder();
int N = Integer.parseInt(st.nextToken());
int M = Integer.parseInt(st.nextToken());

 

다음으로, 배열의 원소들을 입력 받는다.

String inputArr = br.readLine();
st = new StringTokenizer(inputArr);

int[] numArr = new int[N];
for (int i = 0; i < numArr.length; i++) {
    numArr[i] = Integer.parseInt(st.nextToken());
}

 

배열의 각 구간 합을 빠르게 계산하기 위해 Prefix Sum 배열을 만든다.

int[] prefixSum = new int[N + 1];
for (int i = 1; i <= N; i++) {
    prefixSum[i] = numArr[i - 1] + prefixSum[i-1];
}

 

구간 합을 계산하고, 결과를 저장한다.

for(int i = 0; i < M ; ++i){
    String inputTwo = br.readLine();
    st = new StringTokenizer(inputTwo);
    int start = Integer.parseInt(st.nextToken());
    int end = Integer.parseInt(st.nextToken());
    int sum = prefixSum[end] - prefixSum[start - 1];

    sb.append(sum).append("\n");
}

• 각 계산마다 start와 end 값을 읽어온다.
• prefixSum[end] - prefixSum[start - 1]을 계산하여 start부터 end까지의 구간 합을 구한다.
• 결과를 StringBuilder에 저장한다.

 

예를 들어, 입력 배열이 다음과 같다고 가정하자.

numArr = [5, 4, 3, 2, 1]

 

구간 합을 저장하는 배열은 아래와 같이 계산된다.

prefixSum[0] = 0
prefixSum[1] = 5 + 0 = 5
prefixSum[2] = 4 + 5 = 9
prefixSum[3] = 3 + 9 = 12
prefixSum[4] = 2 + 12 = 14
prefixSum[5] = 1 + 14 = 15

 

결과적으로, 구간 합 배열은 아래와 같다.

prefixSum = [0, 5, 9, 12, 14, 15]

 

따라서 문제의 예제를 다시 보면 아래와 같은 계산과정이 도출된다.

start = 1, end = 3

int sum1 = prefixSum[3] - prefixSum[0];
//sum = 12 - 0 = 12

 

start = 2. end = 4

int sum2 = prefixSum[4] - prefixSum[1];
//sum = 14 - 5 = 9

 

start = 5. end = 5

int sum2 = prefixSum[5] - prefixSum[4];
//sum = 15 - 14 = 1

 

현재 구간 합으로 구현한 코드의 시간 복잡도는 O(N + M)이다.

구간 합을 사용하지 않는다면 시간 복잡도는 O(N * M)이 된다.