[BOJ] 10986 나머지 합

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

10986번: 나머지 합

---

문제

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

즉, Ai + ... + Aj (i ≤ j) 의 합이 M으로 나누어 떨어지는 (i, j) 쌍의 개수를 구해야 한다.

---

입력

첫째 줄에 N과 M이 주어진다. (1 ≤ N ≤ 10^6, 2 ≤ M ≤ 10^3)

둘째 줄에 N개의 수 A1, A2, ..., AN이 주어진다. (0 ≤ Ai ≤ 10^9)

---

출력

첫째 줄에 연속된 부분 구간의 합이 M으로 나누어 떨어지는 구간의 개수를 출력한다.

---

풀이

누적 합에 대해 더 자세하게 알게되는 문제이다. 아래는 예시 1과 답을 구하는 코드이다.

void solution() {
    for(int i=0;i<=N;i++){
        ans+=cnt[psum[i]]++;
    }
    cout<<ans;
}

1차원 배열 cnt[n]에는 0~i 까지 나머지가 n인 개수를 저장하고 있다.

 

그렇다면 왜 ans+=cnt[psum[i]]++ 만으로 모든 구간의 개수를 구할 수 있을까?

위 그림은 i = 4일 때, cnt 배열에 들어간 값을 표현 한 것이다. (i = 0 부터 이므로 cnt[0]=3이다)

여기서 cnt[1] = 1 이므로 답의 개수에 +1을 하게 되는데 이는 1,2,3,1에서 앞에 1을 빼고 2,3,1 인 경우를 포함한 것이다.

즉, i = 1이고 j = 4라고 할 때 i+1~j는 3으로 나눠지는 구간이 된다는 것을 알 수 있고 현재 j도 나중에 i가 될 수 있다는 것을 알 수 있다.

따라서 답에 cnt[1] 만큼 더해주고 cnt[1]에 +1을 해줘야하는 것이다.

---

코드

#include <string>
#include <stdio.h>
#include <math.h>
#include <iostream>
#include <vector>
#include <algorithm>
#include <map>
#include <queue>
#include <stack>
#include <climits>
#include <set>

#define p(a, b) make_pair(a, b)
#define SWAP(a, b, type) do { \
    type temp; \
    temp = a;  \
    a = b;     \
    b = temp;  \
} while (0)
#define INF 1000000000
#define endl '\n'
#define ll long long

using namespace std;

ll N,M;
ll psum[1000001];
ll cnt[1001];
ll ans=0;

void input() {
    cin>>N>>M;
    for(int i=1;i<=N;i++){
        cin>>psum[i];
        psum[i]+=psum[i-1];
        psum[i]%=M;
    }
}

void init() {}

void solution() {
    for(int i=0;i<=N;i++){
        ans+=cnt[psum[i]]++;
    }
    cout<<ans;
}

int main()
{
    ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
    input();
    solution();
    return 0;
}