[자바, 코틀린] 백준 17390 - 이건 꼭 풀어야 해! (java, kotlin)

문제 : boj17390

---

 

필요 알고리즘 개념

•  누적 합 (prefix sum)
  • 연속된 범위의 합을 O(1)로 구하기 위해 누적 합을 사용한다. 누적 합에 대한 개념이 있어야 풀 수 있다.
• 정렬
  
  • 정렬이 무엇인지, 자바나 코틀린으로 정렬은 어떻게 하는지 알아야 풀 수 있다.

※ 제 코드에서 왜 main 함수에 로직을 직접 작성하지 않았는지, 왜 Scanner를 쓰지 않고 BufferedReader를 사용했는지 등에 대해서는 '자바로 백준 풀 때의 팁 및 주의점' 글을 참고해주세요. 백준을 자바로 풀어보려고 시작하시는 분이나, 백준에서 자바로 풀 때의 팁을 원하시는 분들도 보시는걸 추천드립니다.

 

---

 

풀이

 

1. 비내림차순으로 정렬해야 한다.

  비내림차순이 생소할 수 있다. 그냥 이 문제에서는 오름차순이라고 생각하면 된다. 자바, 코틀린 모두 내장된 sort 함수를 사용해주면 약 O(NlogN)으로 비내림차순으로 정렬된 배열을 얻을 수 있다. 자바의 경우 Arrays.sort(arr), 코틀린의 경우 arr.sort()를 해주면 된다. 내장 함수를 사용하면 되므로 큰 문제가 되지 않는다. 혹시 내장함수를 쓰지 않고 직접 구현하고 싶다면 퀵정렬이나 병합 정렬로 구현해보자.

 

---

 

2. Q개의 쿼리에 대해 구간합을 구할 수 있어야 한다.

  일반적으로 L부터 R까지의 구간합을 구하려면 L부터 R까지 직접 더해보는 방법이 있다. 총 Q개의 쿼리 이므로 O(Q*(R-L))의 시간복잡도가 필요한데, 어차피 R-L은 최악의 경우 N일 것이므로 O(QN)이 필요하다. 정렬까지 더하면 O(NlogN + QN)으로 Q와 N이 둘다 300,000까지이므로 300,000^2 번 진행해야한다. 당연히 시간 초과로 통과할 수 없다.

 

  전처리를 통해 구간합 배열을 만들어보자. 구간합 배열이란 기존 배열 arr에 대해 prefixSum[i] = arr[0]+arr[1]+...+arr[i] 인 배열을 의미한다. 예를들어 보자면 다음과 같다.

  위와 같이 prefix sum(누적합)을 구해두면 O(1)로 원하는 구간의 합을 알 수 있다. 예를들어 L=1, R=4일 경우 원래는 5+3+2+0을 직접 해봐야 알 수 있다. 하지만 prefixSum배열에서는 prefixSum[R] - prefixSum[L-1] = prefixSum[4] - prefixSum[0] = 10 - 0 = 10 으로 바로 알 수 있다. 왜냐하면 prefixSum[R] = arr[R]+arr[R-1]+...+arr[0]이고, prefixSum[L-1] = arr[L-1] + arr[L-2] + ... + arr[0] 이므로 prefixSum[R] - prefixSum[L-1] = arr[R]+arr[R-1]+...+arr[L+1]+arr[L] 이기 때문이다. 즉 L부터 R까지를 모두 더한 값이 된다.

 

  그럼 Q개에 대해 각각 O(1)로 구간합을 구할 수 있게 되었으니, 최종적으로 시간복잡도는 O(NlogN + Q)가 되어 문제없이 통과 가능하다.  

 

---

 

코드(java) : github

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

public class Main {
    private int getRangedSum(int[] arr, int l, int r) {
        return arr[r] - arr[l-1];
    }

    private void solution() throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        int n = Integer.parseInt(st.nextToken());
        int q = Integer.parseInt(st.nextToken());
        int[] arr = new int[n+1];
        st = new StringTokenizer(br.readLine());
        for (int i = 1; i <= n; i++) {
            arr[i] = Integer.parseInt(st.nextToken());
        }
        Arrays.sort(arr);

        for (int i = 1; i <= n; i++) {
            arr[i] += arr[i-1];
        }

        StringBuilder answer = new StringBuilder();
        while (q-->0) {
            st = new StringTokenizer(br.readLine());
            int l = Integer.parseInt(st.nextToken());
            int r = Integer.parseInt(st.nextToken());
            answer.append(getRangedSum(arr, l, r)).append('\n');
        }
        System.out.print(answer);
    }
    public static void main(String[] args) throws Exception {
        new Main().solution();
    }
}

 

 

코드(kotlin) : github

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

private fun getRangedSum(arr: IntArray, l: Int, r: Int): Int {
    return arr[r] - arr[l-1];
}

fun main() = with(BufferedReader(InputStreamReader(System.`in`))) {
    var st = StringTokenizer(readLine())
    val n = st.nextToken().toInt()
    val q = st.nextToken().toInt()
    var arr = IntArray(n+1)
    st = StringTokenizer(readLine())
    for (i in 1 .. n) {
        arr[i] = st.nextToken().toInt()
    }
    arr.sort()

    for (i in 1 .. n) {
        arr[i] += arr[i-1];
    }

    val answer = StringBuilder();
    repeat(q) {
        st = StringTokenizer(readLine())
        val l = st.nextToken().toInt()
        val r = st.nextToken().toInt()
        answer.append(getRangedSum(arr, l, r)).append('\n')
    }
    print(answer)
}