[백준] 10986번 나머지 합 - C/C++

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

---

 구간합을 구하는 Prefix-sum 유형의 문제이다.

1~i 번째까지 구간의 합을 sum[i] 에 저장을 한 다음 mod_cnt [ ] 배열을 업데이트한다. mod_cnt [i] 배열은 M으로 나눴을 때 나머지가 i가 되는 구간의 개수로 sum [i]를 구한 다음 mod_cnt [ sum [i] % M]을 통해 구할 수 있다.

이제 mod의 특성을 생각해보면 

(A + B) % C = ((A%C) + (B% C))% C로 나머지 연산을 중간에 해도 결과에는 지장이 없다. 즉 i+1부터 j의 구간의 합인 (sum [j] - sum [i]) 를 M으로 나눈 값이 0이라는 것은

(sum [j] - sum [i]) % M =0

sum [j]% M - sum [i]% M = 0;

sum [j] % M = sum [i] % M의 식이 나오게 된다.

위 식으로 알 수 있는 것은 문제인 <연속된 부분 구간의 합이 M으로 나누어 떨어지는 구간>을 구하는 것은 <1~k 번째까지 구간의 합을 M으로 나눴을 때 그 숫자가 같은 쌍들의 수 >가 된다. 이때 추가로 이미 나머지가 0 인 집합들은 굳이 쌍이 아니어도 나머지가 이미 0 이기 때문에 mod_cnt [0]은 한번 더 더해줘야 한다. 쌍들의 수를 구하는 방법은 mod_cnt []가 n 개가 있다면 여기서 nC2를 하면 쌍을 구할 수 있으므로 n*(n-1) / 2로 구할 수 있다.

위에서 설명한 것들을 코드로 나타내면 아래와 같다.

#include <iostream>

using namespace std;

int N, M;
long long mod_cnt[1001] = { 0 };
long long sum[1000001] = { 0 };
int main() {
	cin >> N >> M;
	for (int i = 1; i <= N; i++) {
		cin >> sum[i];
		sum[i] += sum[i - 1];
		mod_cnt[sum[i] % M]++;
	}
	long long ans = mod_cnt[0];
	for (int i = 0; i < M; i++) {
		ans += mod_cnt[i] * (mod_cnt[i] - 1) / 2;
	}
	cout << ans;
}

 

---

세그먼트 트리를 공부하던 친구가 추천해준 문제라 세그먼트 트리로 풀 수 있을 줄 알고 세그먼트 트리로 풀어보니 시간초과가 떴다. 세그먼트 트리를 생성하는 것은 O(logN) 구간 합을 구하는 것도 O(logN)이라고 배워서 시간 안에 풀 수 있을 줄 알았다. 하지만 간과한 것은 한 구간의 합을 구하는 것이 O(logN)이고 모든 구간의 합을 구하려면 (N*N)/2 여서 총 O(N*N*logN)의 시간이 걸리게 된다는 것을 알게 되었다. 그래서 구글링을 한 결과 세그먼트 트리를 이용하는 것보다 단순 배열을 사용하는 것이 이 경우 더 빠르다는 것을 알게 되었다. 이 방법으로는 sum과 mod_cnt배열을 만드는데 O(N), 쌍을 구하는데 O(M)이 걸려 총 O(N+M)의 시간이 걸리게 된다. 

세그먼트 트리를 처음 사용해봐서 시간복잡도를 생각하지 않고 일일이 구간합을 구하는 것보다 빠른 방법이라는 것만 알고 구간합을 보면 세그먼트 트리로 풀어야 한다고만 생각했다. 하지만 세그먼트 트리는 이전 포스팅 같이 한번 한번 구하는 것은 빠르지만 모든 구간의 합을 구하는 문제는 단순 배열을 활용하는 것이 더 빠를 수 도 있다는 것을 알게 되었다. 그래서 세그먼트 트리는 한 번만 구할 때만 좋고 모두 다 구하려면 안 좋은 알고리즘이면 딱히 쓸 수 있는 알고리즘이 아니지 않을까라고 생각을 했는데 원소의 값을 바꾸는 update 가 있다면 단순 배열을 이용한 update는 O(N)의 시간이 걸리지만 세그먼트 트리를 이용한 방법은 O(logN)의 시간이 걸리게 되므로 원소의 값이 자주 바뀌면 세그먼트 트리, 그렇지 않으면 단순 배열로 푸는 것이 효과적일 것 같다.

 

[단순 배열]

- 생성 (1~i번째 까지의 합)   : O(N)

- 특정 구간합 구하기          : O(1)

- 특정 원소 Update            : O(N)

[세그먼트 트리]

- 생성 (여러 구간합의 정보)  : O(logN)

- 특정 구간합 구하기          : O(logN)

- 특정 원소 Update            : O(logN)