[백준] 17390번 이건 꼭 풀어야 해!(누적합 - prefix sum)[Kotlin - 코틀린]

백준 17390번 문제입니다. (solved.ac)기준 실버 3 문제입니다.

https://www.acmicpc.net/problem/17390

17390번: 이건 꼭 풀어야 해!

문제 

문제 접근

첫 번째 줄에 N, Q를 입력 받습니다. N과 Q는 각각 수열 A의 크기와 질문의 개수입니다.

두 번째 줄에는 A배열의 원소들을 입력 받고 그 아래 Q개의 줄에서 질문들을 입력받습니다.

질문은 1 5와 같은 형식인데 A를 비내림차순으로 정리하고(B라는 배열이라고 가정) 거기서 1번째 ~ 5번째까지 원소의 합을 구하는 것입니다.(인덱스 [1]~ [5]가 아님!!!)

저는 아예 두 번째 줄에서 원소를 입력받을 때 정렬하여 B배열을 생성하였습니다.

예제 아래에서 힌트를 볼 수 있는데 힌트에서 설명한 대로 질문이 들어올 때마다 반복문을 통해 반복문을 통해 B라는 배열을 탐색하여 범위의 합을 구한다면 시간 초과 오류가 발생합니다.

저는 누적합을 이용하여 문제를 해결하였습니다. 

누적합을 기록할 sumOfB변수를 만들고 arrayB에서 forEachIndexed 구문을 사용하여 sumOfB에 누적합을 담아 주었습니다. sumOfB는 인덱스 관리를 편하게 하도록 arrayB보다 크기가 하나 더 큰 배열이며(1부터 시작할 것임)

sumOfB[4]에는 arrayB[0] ~ arrayB[3]의 합이 담기게 됩니다.

누적합을 구한 후 질문이 1 ~ 5라면 sumOfB[5] - sumOfB[0]( 5까지의 누적합 - 0까지의 누적합 )을 통해 계산하도록 구현하였습니다.

 

정답 코드

import java.io.BufferedReader
import java.io.BufferedWriter
import java.io.InputStreamReader
import java.io.OutputStreamWriter

// 누적 합.
fun main() {

    val br = BufferedReader(InputStreamReader(System.`in`))
    val bw = BufferedWriter(OutputStreamWriter(System.out))

    val (N, Q) = br.readLine().split(" ").map { it.toInt() }
    var arrayB = br.readLine().split(" ").map { it.toInt() }.toIntArray().sortedArray()
    var sumOfB = IntArray(arrayB.size + 1, { 0 }) // 인덱스 계산을 편하게 하기 위해.

    // 정렬 후 누적 합 계산. ex) sumOfB[3]은 arrayB[0] ~ arrayB[2]의 합이 담기게 됨.
    arrayB.forEachIndexed { index, it ->
        sumOfB[index + 1] = sumOfB[index] + it

    }

    repeat(Q) {
        val range = br.readLine().split(" ").map { it.toInt() }
        var ans = sumOfB[range[1]] - sumOfB[range[0] - 1]

        bw.write("$ans\n")

    }
    bw.flush()
    bw.close()

}