10986번: 나머지 합

10986번: 나머지 합

사용 언어: C++

문제 요약

수 N개 A1, A2, ..., AN이 주어진다. 이때 연속된 부분 구간의 합이 M으로 나누어 떨어지는 구간의 개수를 구하는 프로그램을 작성하시오.

풀이

처음에는 아래와 같이 모든 (i, j) 쌍에 대해 부분합을 구하고,

그 값이 M으로 나누어떨어지는지 검사하려 했지만 시간 초과가 났다.

for (int i = 1; i <= N; i++) 
{
    for (int j = i; j <= N; j++) 
    {
        if ((prefix[j] - prefix[i-1]) % M == 0)
            ans++;
    }
}

그래서 시간을 어떻게 줄일 수 있을까 고민해봤는데, 위의 O(N^2)이 아닌 O(N)으로 줄일 수 있었다.

방법은 누적합 + 나머지 활용이다. 누적합을 다 하고, 나머지가 같은 prefix 값들끼리 nC2 로 조합하는 것이다.

예를 들어 나머지가 1인 prefix 수가 4개면, 이 4개 중 2개 씩 고르면 prefix[4] - prefix[2]를 하면 얘네도 결국 M으로 나누어 떨어지므로 저 구간도 포함되는 것이다. 아래가 그 코드다.

#include <iostream>
#include <vector>

using namespace std;

int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);

    int N, M;
    cin >> N >> M;

    vector<long long> prefix(N+1, 0);
    vector<long long> cnt(M, 0);     
    long long ans = 0;

    for (int i = 1; i <= N; i++) 
    {
        int x;
        cin >> x;
        prefix[i] = prefix[i - 1] + x;
        int r = prefix[i] % M;
        if (r == 0) ans++;
        cnt[r]++;
    }

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

    cout << ans << '\n';
    return 0;
}