백준 10986 나머지 합 Kotlin (누적 합)

문제 출처 : https://www.acmicpc.net/problem/10986

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으로 나누어 떨어지는 구간의 개수를 출력한다.

알고리즘 분류

• 수학
• 누적 합

풀이

누적 합인데 나머지 연산의 성질을 이용하는 문제라서 좀 어려웠다.

나머지의 성질에 의하면 우선

(A+B) % M이 0이라면

(A%M + B%M) % M 도 0이다.

즉, (A+B) % M 을 (A%M + B%M) % M으로 나타낼 수 있다는 점을 이용해야 한다.

이 성질은 마이너스, 곱하기에 대해서도 동일하게 적용할 수 있다.

 

 

우선 위의 예시에 대한 누적 합을 구한다.

preSum = 1 3 6 7 9

위의 공식에 따라 이 누적 합에 %m을 하면

preSum = 1 0 0 1 0 이 된다.

우선 여기서 m으로 나누어 떨어지는 구간의 합이 3개가 나왔음을 알 수 있다.

 

추가로 preSum이라는 누적 합을 이용한다면 어떠한 구간 (i, j)의 합은 preSum[j] - preSum[i-1]이다.

우리는 (preSum[j] - preSum[i-1]) %m == 0인 경우를 모두 찾아야 하는데,

preSum[j] % m = preSum[i-1] % m 인 쌍의 경우를 모두 찾으면 된다.  

 

 

preSum[x] % m이 1인 경우가 2개, 

preSum[x] % m이 0인 경우가 3개로,

나머지 값이 1일 때의 쌍은 2C2

나머지 값이 0일 때의 쌍은 3C2

즉 preSum[i] % m = preSum[j-1] % m 인 쌍의 경우는 총 4가지다.

 

코드

import java.util.*

val br = System.`in`.bufferedReader()

fun main() = with(System.out.bufferedWriter()){

    val (n,m) = br.readLine().split(' ').map{it.toInt()}
    val tk = StringTokenizer(br.readLine())
    val preSum = LongArray(n+1)
    var idx=1
    val cnt = LongArray(m)

    while(tk.hasMoreTokens()){
        preSum[idx] = tk.nextToken().toLong()
        preSum[idx] = (preSum[idx]+preSum[idx-1])%m
        cnt[preSum[idx++].toInt()]++
    }
    var answer = cnt[0]

    for(i in 0 until m){
        answer += cnt[i] * (cnt[i]-1) /2
    }
    write("$answer")
    close()
}