BOJ-10986 나머지 합 Swift

문제

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

내가 푼 풀이

첫번째 접근방법: 세그먼트 트리(시간초과)

구간합트리로 해당 구간합을 구해 나누어 떨어지는지 확인해보았다.

구간합트리의 모든노드는 모든 구간합의 경우의수가 아니였고, 구간합트리를 만들어도 모든 구간합의 경우의수를 구하려니 2중 반복문으로 시간초과가 나게되었다.

 

이 문제는 수학적지식을 이용해서 풀어야만 했었다. 지식의 한계를 느끼고 다른사람의 해석을 보았다.

 

두번째 접근방법: 합배열과 구간합

구간합배열은 이렇게 만들 수 있다.

S(1), S(2), S(3)....

S(1)은 1부터1까지, S(2)는 1부터2까지, S(3)은 1부터3까지의 구간합

	A[1]	A[2]	A[3]	A[4]	A[5]
A	1	2	3	1	2
S	1	3	6	7	9
S[i] % M	1	0	0	1	0

 

여기서 구간합배열인 S를 M으로 나누어보면 아래 배열과 같다 [1, 0, 0, 1, 0]

이를 해석해보면 구간합 [1~2], [1~3], [1~5]는 M과 나누어떨어지고, [1~1], [1~4]는 나머지 1이 남는다.

또한 구간합은 S[i~j] = S[j] - S[i-1]로 나타낼 수 있다. 예로 2에서4까지의 구간합은 S[4] - S[1]이다.

이 특징을 이용하면 S[4] - S[1] = 1 - 1 = 0이되고 이는 M으로 나누어떨어진다고 볼 수 있다.

따라서 우리는 아래처럼 종합할 수 있다.

모든 구간합의 경우의수는 아래와 같다.

1. 구간합배열S를 M으로 나누었을때의 0이되는 구간합

2. 나머지배열중 서로 나머지가 같아서 S[i~j] = S[j] - S[i-1] = 0이되는 두개의 나머지배열

 

1번의 경우의수는 S[2], S[3], S[5] 총 3개가 된다.

2번의 경우의 수는 같은 나머지원소중에서 2개를 뽑으면된다(S[j] - S[i-1] = 0을 성립하기 위한 S[j]와 S[i-1]이 필요하다.)

나머지가 0인것중에서 2개를 뽑는 수: ${_3}C{_2}$ -> 3*(3-1)/2 = 3

나머지가 1인것중에서 2개를 뽑는 수: ${_2}C{_2}$  -> 2*(2-1)/2 = 1

3 + 3 + 1 = 7 이 된다.

 

M으로 나누었을때 나머지는 0부터 M-1까지 있으므로 모두 순회하며 더해준다.

 

코드로 구현하면 다음과 같다.

import Foundation

// 입력받기
let input = readLine()!.split(separator: " ").map{Int(String($0))!}
let N = input[0], M = input[1]
// 주어진배열
var arr = readLine()!.split(separator: " ").map{Int(String($0))!}
// 합배열
var S = [Int]()
var sum = 0
// 합배열을 M으로 나눈 나머지배열
var divided = [Int]()
// 나누어떨어지는 갯수
var count = 0

// 합배열 만들기
for i in arr {
    sum += i
    S.append(sum)
}

// 나머지 배열만들기
divided = S.map{$0 % M}
// 나머지 배열중 나머지가 0이면 그 구간합은 나누어떨어지므로 더해준다.
count = divided.filter{$0 == 0}.count

// 나머지배열중 같은 나머지의 갯수중에 2개를 고르는 경우의수를 더한다.
// S[i~j] = S[j] - S[i-1] = 0이 되는 두개의 구간합이 필요하므로
for i in 0..<M {
    let num = divided.filter{$0 == i}.count
    count = count + (num * (num-1))/2
}
print(count)