[백준] 1644번 소수의 연속합 (Kotlin)

문제

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

문제 요약

• 자연수가 주어졌을 때 연속된 소수의 합으로 나타낼 수 있는 경우의 수는 몇 가지인가?
• 시간제한 : 2초
• input : 1 이상 4,000,000 이하의 자연수

접근

• 소수를 구하기 위해 에라토스테네스의 체를 사용
• 범위내에 있는 합을 빠르게 구하기 위해 구간 합 사용
• 시작과 끝의 범위를 투 포인터 알고리즘 사용

구현

• 입력값 이하의 소수 구하기

private fun getPrimeNumber(n: Int): ArrayList<Int> {

    val primeNumber = arrayListOf<Int>()
    val isPrime = BooleanArray(size = n + 1) { true }

    if (n == 1) return primeNumber

    isPrime[0] = false
    isPrime[1] = false

    for (i in 2..n) {
        if (!isPrime[i]) continue
        else {
            primeNumber.add(i)
            for (j in i * 2..n step (i)) {
                isPrime[j] = false
            }
        }
    }

    return primeNumber
}

• 구간 합 구하기

private fun getPrefixSum(primeNumber: ArrayList<Int>): ArrayList<Int> {

    val prefixSum = arrayListOf<Int>()

    prefixSum.add(0)
    for (i in primeNumber.indices) {
        prefixSum.add(prefixSum[i] + primeNumber[i])
    }

    return prefixSum
}

• 결과 (연속된 소수의 합으로 나타낼 수 있는 경우의 수) 구하기

private fun getResult(n: Int, prefixSum: ArrayList<Int>): Int {

    var result = 0
    var leftPointer = 0
    var rightPointer = 0

    while (leftPointer < prefixSum.size && rightPointer < prefixSum.size) {

        val sum = prefixSum[rightPointer] - prefixSum[leftPointer]
        if (sum == n) result++

        if (sum < n) rightPointer++
        else leftPointer++
    }

    return result
}

최종 코드

/*
* 백준 1644번. 소수의 연속합
* https://www.acmicpc.net/problem/1644
*/

private fun main() {
    val n = readln().toInt()
    val result = getResult(n)
    printResult(result)
}

private fun getResult(n: Int): Int {
    val primeNumber = getPrimeNumber(n)
    val prefixSum = getPrefixSum(primeNumber)
    return getResult(n, prefixSum)
}

private fun getPrimeNumber(n: Int): ArrayList<Int> {

    val primeNumber = arrayListOf<Int>()
    val isPrime = BooleanArray(size = n + 1) { true }

    if (n == 1) return primeNumber

    isPrime[0] = false
    isPrime[1] = false

    for (i in 2..n) {
        if (!isPrime[i]) continue
        else {
            primeNumber.add(i)
            for (j in i * 2..n step (i)) {
                isPrime[j] = false
            }
        }
    }

    return primeNumber
}

private fun getPrefixSum(primeNumber: ArrayList<Int>): ArrayList<Int> {

    val prefixSum = arrayListOf<Int>()

    prefixSum.add(0)
    for (i in primeNumber.indices) {
        prefixSum.add(prefixSum[i] + primeNumber[i])
    }

    return prefixSum
}

private fun getResult(n: Int, prefixSum: ArrayList<Int>): Int {

    var result = 0
    var leftPointer = 0
    var rightPointer = 0

    while (leftPointer < prefixSum.size && rightPointer < prefixSum.size) {

        val sum = prefixSum[rightPointer] - prefixSum[leftPointer]
        if (sum == n) result++

        if (sum < n) rightPointer++
        else leftPointer++
    }

    return result
}

private fun printResult(result: Int) {
    val bw = System.out.bufferedWriter()
    bw.write("$result\n")
    bw.flush()
    bw.close()
}