[백준] 수들의 합2 - 실버 4 (Kotlin)

문제

N개의 수로 된 수열 A[1], A[2], …, A[N] 이 있다. 이 수열의 i번째 수부터 j번째 수까지의 합 A[i] + A[i+1] + … + A[j-1] + A[j]가 M이 되는 경우의 수를 구하는 프로그램을 작성하시오.

입력

첫째 줄에 N(1 ≤ N ≤ 10,000), M(1 ≤ M ≤ 300,000,000)이 주어진다. 다음 줄에는 A[1], A[2], …, A[N]이 공백으로 분리되어 주어진다. 각각의 A[x]는 30,000을 넘지 않는 자연수이다.

출력

첫째 줄에 경우의 수를 출력한다.

풀이

문제에서 주어진 조건 N 의 범위가 1 ~ 10,000 까지 이므로 무작정 이중포문을 사용하여 문제를 푼다면 O(N^2) 하여 시간초과가 날 가능성이 있다.
따라서, 투포인터 알고리즘을 사용하면 아주 간단하게 풀 수 있는 문제이기 때문에 투포인터를 직접 구현해보는 연습을 했다.
투포인터 알고리즘을 문제에 적용해보면

• 시작 점과 끝 점이 첫 번째 원소의 인덱스를 가리키도록 한다.
• 현재 부분 합이 M과 같다면 카운트한다.
• 현재 부분 합이 M보다 작거나 같다면 끝 점을 1 증가시킨다.
• M보다 크다면 시작 점을 1 증가시킨다.
• 모든 경우를 확인할 때까지 2 ~ 4 과정을 반복한다.

소스 코드

fun main(args: Array<String>) {
    /*
4 2
1 1 1 1
     */
    var start = 0
    var end = 0
    var count = 0
    var sum = 0
    val list = ArrayList<Int>()

    val (n, m) = readLine()!!.split(' ').map { it.toInt() }
    readLine()!!.split(' ').forEach { list.add(it.toInt()) }

    while (start != n && end != n) {
        for (i in start..end) {
            sum += list[i]
        }
        if (sum == m) {
            count++
            end++
        }
        else if (sum < m) end++
        else if (sum > m) start++
        sum = 0
    }

    println(count)
}

소스 코드(재업로드)

처음에 생각한 위의 코드는 다시 생각해보니

for (i in start..end) {
	...
}

이 부분 때문에 반복문을 여러 번 돌리는 방식이랑 시간 복잡도가 다를게 없었다. (바보)
따라서 아래와 같이 코드를 변경했다.

fun main(args: Array<String>) {
    /*
4 2
1 1 1 1
     */
    var start = 0
    var end = 0
    var count = 0
    var sum = 0
    val list = ArrayList<Int>()

    val (n, m) = readLine()!!.split(' ').map { it.toInt() }
    readLine()!!.split(' ').forEach { list.add(it.toInt()) }

        for (start in 0 until n) {
            while (sum < m && end < n) {
                sum += list[end++]
            }

            if (sum == m) {
                count++
            }
            sum -= list[start]
        }

    println(count)
}