수 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으로 나누어 떨어지는 구간의 개수를 출력한다.
예제 입력 1
5 3
1 2 3 1 2
예제 출력 1
7
누적합 문제는 preSum 구현하는 방식을 기본으로 하여, 이를 어떻게 활용할 것인가가 핵심이다.
문제에서 N이 최대 10^6이기에, 시간복잡도 상 O(NlogN)이하가 나와야한다는 것을 파악한 상태로 문제에 접근하고 있었다.
문제에서의 핵심은 나머지가 0이 되는 구간합을 찾는 것이기에, 나머지에 맞추어 구분해야 한다는 사고까지는 하였으나, 그 이후 과정은 생각하지 못하고 멈춰 있었다.
게시판에 도움 글을 통해 갈피를 잡고 구현할 수 있었다.
즉, 두 부분합의 차이가 특정 구간의 합이 될 수 있다는 것이 핵심이다.
그래서 나머지가 같은 부분합끼리 묶어준 후, 나머지가 같은 그룹에서 부분합 두 개를 뽑아 이 둘을 서로 빼주면, 나머지가 0이 될 수 있다는 것이다.
이렇게 알고리즘을 파악하고 구현하는 와중, 예제에 답이 계속 4가 나왔었다.
그 이유는
result += remain[0];
를 하지 않았기 때문이였다. 처음에는 두 부분합을 뽑아 둘을 서로 빼주는 경우의 수는 그룹의 사이즈가 n일 때, nC2와 동일하기에 n * (n - 1) / 2만 생각하고 있었는데, 애초에 나머지가 0인 경우도 하나의 경우의 수가 될 수 있다는 것을 배제하고 있었다.
위와 같은 부분을 조심해야 할 것 같다.
#include <bits/stdc++.h>
using namespace std;
#define endl "\n"
#define ll long long
int n, m;
ll preSum[1000001];
ll remain[1001];
ll solve();
int main() {
ios::sync_with_stdio(false);
cin.tie(NULL);
cout.tie(NULL);
memset(preSum, 0, sizeof(preSum));
cin >> n >> m;
for(int i = 1; i < n + 1; i++) {
cin >> preSum[i];
preSum[i] += preSum[i - 1];
remain[preSum[i] % m]++;
}
cout << solve() << endl;
return 0;
}
ll solve() {
ll result = 0;
result += remain[0];
for(int i = 0; i < m; i++) {
if(remain[i] >= 2) {
result += (remain[i] * (remain[i] - 1)) / 2;
}
}
return result;
}