[Baekjoon 10986 | C++] Remainder Sum

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

다음과 같이 코드를 작성하게 되면, 시간초과로 실패하게 된다.

#include <iostream>
#include <vector>
using namespace std;

vector<int> v;
int mem[1'000'000];

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

    int N, M;
    cin >> N >> M;
    int input;
    for (int i = 0; i < N; ++i) {
        cin >> input;
        v.push_back(input);
    }

    int result = 0;
    for (int j = 0; j < N - 1; ++j) {
        mem[j] = v[j];
        if (mem[j] % M == 0) {
            result++;
        }
        for (int i = j; i < N - 1; ++i) {
            mem[i + 1] = mem[i] + v[i + 1];
            if (mem[i + 1] % M == 0) {
                result++;
            }
        }
    }
    cout << result;

    return 0;
}

그렇기 때문에 새로운 방식인 누적합으로 문제를 해결할 수 있다.

문제에서 제공하는 기본 예제를 통해 이해해보자.

Input

5 3
1 2 3 1 2

Output

7

누적합으로 나타내면 다음과 같다.

여기서 누적합%M 부분을 자세히 보아야 한다.

예를 들어 누적합 나머지 값이 1인 값들인 1과 7의 경우, 7 - 1 = 6으로 나머지 값이 0이 된다.

따라서 누적합 나머지 값이 0인 경우에서 두 값을 빼지 않을 경우 3가지,

나머지 값이 0인데 두 값을 뺄 경우는 3C2 = 3(가지),

나머지 값이 1이고 두 값을 뺄 경우는 2C2 = 1(가지),

3 + 3 + 1 = 7(가지) 가 나오게 된다.

이를 적용한 코드는 다음과 같다.

#include <iostream>
using namespace std;

using ll = long long;

ll mem[1'001];

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

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

    ll sum = 0;
    for (int i = 0; i < N; ++i) {
        cin >> input;
        sum += input;
        mem[sum % M]++;
    }

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

    cout << result;

    return 0;
}