나머지 합

알고리즘 구분 : 수학, 누적합

문제

수 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;
}

• 문제 출처 : https://www.acmicpc.net/problem/10986
• 해당 풀이는 백준에서 "맞았습니다!!" 판정을 받았습니다.