[PS] BOJ 나머지 합 - 10986 (C++)

문제 링크 - 나머지 합

문제 정리

자연수 N, M 이 주어지고, N개의 수가 주어질때, 연속된 부분 구간의 합이 M 으로 나누어 떨어지는 구간의 갯수를 구하는 문제이다. 단순하게 봤을 땐, Brute Force 로 모든 구간 쌍에 대해 각 구간합을 계산하여 풀면 될거같지만, 당연하게도 이런 방식으로 풀면 time limit 에 걸린다. 이 경우 Time complexity 는 다음과 같다:

• 모든 구간 쌍 combination 찾기: O(n^2)
• 각 구간의 합 계산: O(n)
• 따라서 전체 시간 복잡도는 O(n^3) 이 나온다.

'연속된 부분 구간의 합' 이라는 문구를 보면 어렵지 않게 누적합을 사용해야 한다는 걸 알 수 있었다. 하지만 그 다음에 어떤 방식으로 진행해야 할지 어려움이 있었는데, 여러 고수들의 풀이과정을 보며 해결 전략을 정리해 보았다.

해결 전략

i 번째 index 까지의 누적 합을 p_sum(i)라 할때, [i + 1, j] 부분 구간의 누적합은 p_sum[j] - p_sum[i] 로 표현할 수 있다. [i + 1, j] 구간의 누적 합이 M 으로 나누어 떨어지기 위해선 다음 식이 성립해야 한다.

(p_sum[j] - p_sum[i]) % M = 0

모듈러 연산 공식에 따라 위의 식은 아래와 같이 표현이 가능하다.

p_sum[j] % M - p_sum[i] % M  = 0 -> p_sum[j] % M = p_sum[i] % M

즉, index 까지의 누적 합의 나머지 값이 같은 indices 들의 combinations 갯수를 모두 더하면 된다!

한가지 주의해야 할 edge case 는, 나머지 값이 0인 경우, 처음부터 해당 index 까지의 합도 M 으로 나누어 떨어지기 때문에 나머지 값이 0인 indices 의 수도 추가로 더해줘야 한다.

예시

아래의 예시를 위의 풀이과정을 통해 풀어보자

[1, 2, 3, 1, 2], M = 3

우선, index 까지의 누적합을 구한다.

[1, 3, 6, 7, 9]

그 후, M 으로 나눈 나머지를 저장한다.

[1, 0, 0, 1, 0]

나머지가 0인 값이 3개, 나머지가 1인 값이 2개가 나오는 걸 알 수 있다. 각 경우에 대해 combinations 수를 구한 후 더해주면 된다!

0: 3 -> (3 * (3 - 1) / 2) = 3
1: 2 -> (2 * (2 - 1) / 2) = 1

위의 값들은 [i, j] 범위에 대해서 [i + 1, j] 의 값들에 대한 누적합을 구한 것이므로 index 가 0 인 처음부터 시작하는 누적합에 대한 값을 추가로 더해줘야 한다. 위에서 얘기한 바와 같이, 나머지가 0인 누적합의 수를 모두 더해주면 해결된다.

0: 3 -> (3 * (3 - 1) / 2) = 3
1: 2 -> (2 * (2 - 1) / 2) = 1

3 + 1 + 3(나머지가 0인 누적합의 수) = 7

문제 풀이

C++ 코드로 문제를 풀어보자. 우선 문제를 풀기전에 고려한 내용들을 정리해 보았다.

• 누적 합의 modulo 를 리스트로 만든 후, 리스트를 돌며 각 modulo 의 수를 세지 않고, input 을 받으면서 누적 합과 modulo 를 즉시 계산 한 후 나머지 값을 index 로 하는 리스트의 카운트를 올려주는 방식을 사용.
• 문제상의 숫자 범위가 크므로 경우에 따라 int 가 아닌 long long 자료형을 사용.

나머지 합 변수의 범위는 다음과 같다.

이때, 아래 상황들에서 int 의 범위를 초과할 수 있다.

• 누적 합 (runningSum) 은 최악의 경우 10^6 * 10^9 = 10^15 가 되므로 int 를 쓸 경우 overflow 가 생길 수 있다.
• Combination 계산 과정 (n _ (n - 1) / 2) 에서 n _ (n - 1) 을 계산할 때 최악의 경우 10^12 가 되므로 int 를 쓸 경우 overflow 가 생길 수 있다.

코드

#include <iostream>
using namespace std;
int N, M;
long long moduloCounts[1001];

void solve()
{
    cin >> N >> M;
    long long runningSum = 0;
    long long res = 0;

    for(int i = 1; i <= N; i++)
    {
        int num; cin >> num;
        runningSum += num;
        moduloCounts[runningSum % M]++;
    }

    for(int i = 0; i <= M; i++)
    {
        long long n = moduloCounts[i];
        res += n * (n - 1) / 2;
    }

    cout << res + moduloCounts[0]<< '\n';

}

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

    solve();

    return 0;
}