10986번: 나머지 합 - Kotlin (누적 합)

10986번: 나머지 합

문제

수 N개 A1, A2, ..., AN이 주어진다. 이때, 연속된 부분 구간의 합이 M으로 나누어 떨어지는 구간의 개수를 구하는 프로그램을 작성하시오.

즉, Ai + ... + Aj (i ≤ j) 의 합이 M으로 나누어 떨어지는 (i, j) 쌍의 개수를 구해야 한다.

입력

첫째 줄에 N과 M이 주어진다. (1 ≤ N ≤ 106, 2 ≤ M ≤ 103)

둘째 줄에 N개의 수 A1, A2, ..., AN이 주어진다. (0 ≤ Ai ≤ 109)

출력

첫째 줄에 연속된 부분 구간의 합이 M으로 나누어 떨어지는 구간의 개수를 출력한다.

 

 

나의 풀이 - 1차 (메모리 초과)

 누적합을 이전에 해결해 봤기에 구간들의 합을 2차원 배열로 정의할 수 있겠다는 생각이 들어서 다음과 같이 구현하였다.

 

예) 1행 3열은 1부터 3까지의 누적합

 

하지만 메모리 초과가 발생하였다.

fun main() = with(System.`in`.bufferedReader()) {
    val (n, m) = readLine().split(" ").map { it.toInt() }
    val num = readLine().split(" ").map { it.toInt() }

    val dp = Array(n + 1) { IntArray(n + 1) }


    for (i in 1..n) {
        dp[0][i] = num[i - 1] + dp[0][i - 1]
    }

    var cnt = 0
    for (i in 1..n) {
        for (j in i..n) {
            dp[i][j] = dp[0][j] - dp[0][i - 1]
            if (dp[i][j] % m == 0) cnt++
        }
    }

    println(cnt)
}

 

 

나의 풀이 - 2차

 메모리 초과가 발생한 이유는 2차원 배열 dp의크기가 (n + 1) * (n + 1)으로 설정되어 있기 때문에 N = 10^6인 상황에서는 메모리 용량을 초과하게 된다. 해서 다른 방식을 사용해야 한다.

 

dp에 값을 넣어줄 때, 누적 합을 넣어주지 말고 m의 나머지 값을 넣어주는 방식으로 진행해야 한다. 

 

fun main() = with(System.`in`.bufferedReader()) {
    val (n, m) = readLine().split(" ").map { it.toInt() }
    val num = readLine().split(" ").map { it.toLong() }

    val dp = LongArray(n + 1)
    var cnt = LongArray(m)

    for (i in 1..n) {
        dp[i] = (num[i - 1] + dp[i - 1]) % m
        cnt[dp[i].toInt()]++
    }

    var ans = cnt[0]
    for (i in 0 until m) {
        ans += cnt[i] * (cnt[i] - 1) / 2
    }
    println(ans)
}