10986번 (gold 3)

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

구간합 알고리즘을 사용하는 문제중 골드문제로 어려운 축에 속한다.

 

#include <iostream>
#include <vector>

using namespace std;
int main() {

	long N, mod;
	cin >> N >> mod;
	vector<long> arr_sum(N + 1, 0);
	vector<long> mod_count(mod, 0);
	long count = 0;

	for (int i = 1; i <= N; ++i) {
		long temp;
		cin >> temp;
		arr_sum[i] = arr_sum[i - 1] + temp;
	}

	for (int i = 1; i <= N; ++i) {
		mod_count[arr_sum[i] % mod]++;
		if (arr_sum[i] % mod == 0) {
			count++;
		}
	}
	for (int i = 0; i < mod; ++i) {
		long tmp = mod_count[i];
		count += (tmp * (tmp - 1)) / 2;
	}

	cout << count;
	return 0;
}

 mod 연산과 조합공식을 이해하고 문제를 풀면 조금 더 쉬운데, 구간합에 나머지 연산 즉 모듈러 연산을 적용해서 나머지가 0이면 그 구간이 나누어 떨어지는 곳이기 때문에, 그 만큼 count를 더하고 추가로 각각의 구간합의 나머지 연산이 같다면, 예를들어 예제 입력에서 

5 3

1 2 3 1 2를 입력하면 7이 출력된다는데, 1 2 3 1 2를 컨테이너 안에 구간합을 저장해 놓는다면

arr_sum[1] == 1, arr_sum[2] == 2 + arr_sum[1], arr_sum[3] == 3 + arr_sum[2], arr[4] == 1 + arr_sum[3],

arr[5] == 2 + arr_sum[4]이고, arr_sum[i] % m == 0이라면, 그 구간합은 나누어 떨어지기 때문에 count++가 된다.

 

 또 arr_sum에 모든 요소에 예제에 있는 m 즉 3을 나누면 1 3 6 7 9 --> 1 0 0 1 0 이 되는데 저 다섯구간에서 같은 나머지를 가지는 구간들의 사이 구간합은 3으로 나누어 떨어진다고 생각할 수 있는데, 

(a - b) % m = ((a % m) - (b % m)) % m 이기 때문이다. 이산수학의 지식이 조금 필요함.

 

 즉 1하고 7사이의 구간합은 6이고 6 % 3 = ((1 % 3) - (7 % 3)) % 3 = 0이다. 즉 구간합에 나머지 연산을 한 값이 같은 구간들의 사이 구간합은 나머지 값이 0이 된다!

그래서 1 0 0 1 0 에서 1이 2개, 0이 3개인데, 1이 2개중 2개를 뽑는 경우의 수 2C2와 3개의 0중 2개를 뽑는 경우의 수

3C2를 count에 추가로 더해주면 된다!