백준 10986. 나머지 합(C++)

10986. 나머지 합

10986번: 나머지 합

---

입력

• 첫째 줄에 N과 M이 주어진다.
• 둘째 줄에 N개의 수 A1, A2, ..., AN이 주어진다.

---

문제 풀이

• Ai + ... + Aj 의 합이 M으로 나누어 떨어지는 (i, j) 쌍의 개수를 구해야 한다.
• 가장 먼저 할 수 있는 생각은 다중 for문을 통해 모든 경우의 수를 확인하는 것이다.
• 하지만 위와 같은 방법으로 한다면 N(수의 개수)가 커짐에 따라 매우 많은 시간이 걸리게 된다.
• 여기서 우리는 누적합을 구해 연속된 구간의 합을 빨리 구할 수 있게 도와주는 Prefix 알고리즘을 활용할 수 있다.
• 먼저, 수열을 입력받은 후 각 인덱스에서의 누적합을 구한다. 1중 for문으로 쉽게 구할 수 있다.
• 그렇다면 Sum[i] - Sum[j]를 통해 (i, j)의 구간합을 빠르게 구할 수 있다.
• 여기서 Sum[i]를 M으로 나눴을 때 나머지와 Sum[j]를 M으로 나눴을 때 나머지가 같아야 해당 구간합이 M으로 나누어 떨어진다고 할 수 있다. 따라서 각 나머지별 구간합의 수를 구하고 잘 조합하면 정답을 도출할 수 있다.
• 코드는 다음과 같다.

#include <iostream>
#include <vector>
#include <algorithm>

int N, M;
long long result;
std::vector<long long>arr(1);
std::vector<long long>modarr;

int main()
{
	std::ios_base::sync_with_stdio(false);
	std::cin.tie(NULL);
	std::cout.tie(NULL);

	std::cin >> N >> M;

	for (int i = 0; i < N; i++) {
		long long tmp;
		std::cin >> tmp;
		arr.push_back(tmp);
	}

	for (int i = 0; i < M; i++) {
		modarr.push_back(0);
	}

	for (int i = 1; i <= N; i++) {
		arr[i] = arr[i] + arr[i - 1];  
	}

	for (int i = 1; i <= N; i++) {
		arr[i] = arr[i] % M;
	}

	for (int i = 0; i <= N; i++) {
		modarr[arr[i]] += 1;
	}

	for (int i = 0; i < M; i++) {
		if (modarr[i] >= 2) {
			result += (modarr[i] * (modarr[i] - 1)) / 2;
		}
	}

	std::cout << result;

	return 0;
}