[Swift] 백준 10986 : 나머지 합

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

10986번: 나머지 합

 

문제

 

단순히 누적합 알고리즘으로 접근할 경우에 아래와 같이 시작점과 끝점의 모든 경우의 수를 확인해야 하므로 시간 복잡도가 O(n^2)

for i in 0...n-1 {
	for j in i+1...n {
    	count += ( psum[j]-psum[i] ) % 1 == 0 ? 1 : 0
    {
{

 

 

하지만 여기서 누적합 psum[0]...psum[n] 중 나머지값이 동일할 psum[i], psum[j]이 있다면 부분합 psum[j]-psum[i]의 나머지는 0이라는 원리와 이를 경우의 수로 계산하는 접근 방식으로 시간 복잡도를 O(n+m)으로 줄일 수 있다.

 

누적합 psum을 만들면서 동시에 psum[i]의 나머지가 0...m-1일 빈도수를  surplus[0...m-1]에 입력한다.

for i in 1...n {
    psum[i] = psum[i-1] + inputArr[i-1]
    surplus[psum[i]%m] += 1
}

 

그 다음 나머지가 동일한 누적합들 중 두개를 고르는 경우의 수를 구한다. 즉 (psum[j]-psum[i])%m 이 0이 될 수 있는 경우의 수를 구한다. 먼저 n 개 중 2개를 고르는 경우의 수를 구하는 메서드

func comb (_ n: Int, _ m: Int) -> Int {
    if n < 2 {
        return 0
    }
    return n * (n - 1)/m
}

 

그 다음 이 메서드를 이용해 나머지가 같은 부분합 중에 두개를 선택하는 모든 경우의 수를 구한다. 

var count : Int = surplus[0]

for i in 1..<m {
    count += comb(surplus[i], 2)
}

 

참고로 surplus[0]이 count의 초기값인 이유는 (psum[i]-psum[0])%m = 0 이면 그 누적합 자체로 Arr[1]..부터 Arr[i]까지 더한 부분합이고 동시에 나머지가 0인 부분합이기 때문에 그 개수를 카운트해준다.