백준 10986번 나머지 합 (Java)

문제

백준 10986번 나머지 합 / 골드3
수 N개 A₁, A₂, ..., A_N이 주어진다. 이때, 연속된 부분 구간의 합이 M으로 나누어 떨어지는 구간의 개수를 구하는 프로그램을 작성하시오.
즉, A_i + ... + A_j (i ≤ j) 의 합이 M으로 나누어 떨어지는 (i, j) 쌍의 개수를 구해야 한다.

입력

첫째 줄에 N과 M이 주어진다. (1 ≤ N ≤ 10⁶, 2 ≤ M ≤ 10³)
둘째 줄에 N개의 수 A₁, A₂, ..., A_N이 주어진다. (0 ≤ A_i ≤ 10⁹)

풀이

누적 합 문제이다. 하지만 단순히 모든 구간의 누적 합을 구하게 되면 입력 N이 10⁶이므로 시간복잡도는 O(n²)로 시간 초과가 발생한다.

그래서 M으로 나누었을 때, 나머지가 같은 수를 카운트 해주는 배열 remainders를 만들었다. remainders배열의 index가 나머지이다. 이 나머지가 같은 수 2개를 빼면 M의 배수가 나오게 된다. 그렇기에 각 배열에 저장된 수를 조합 계산하면 remainders[i] = n 일 때, n(n-1)/2 이라는 계산식이 나온다.

누적 합을 구하기 위해서는 이전 누적 합을 저장하는 prevSum과 현재 누적 합을 저장하는 sum, 누적 합이 M으로 나누어떨어지는 수를 저장하는 remainderZero를 사용했다. 이 누적 합들의 수 2개를 빼면 마찬가지로 M의 배수가 나오게 된다. 그렇기에 remainderZero = n 일 때, n+(n(n-1)/2) = n(n+1)/2 이라는 식이 나오게 된다.

코드

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.StringTokenizer;

public class Main {
    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        int N = Integer.parseInt(st.nextToken());
        int M = Integer.parseInt(st.nextToken());

        long sum = 0;
        long prevSum = 0;
        long remainderZero = 0;
        long[] remainders = new long[M];

        st = new StringTokenizer(br.readLine());
        for (int i = 1; i <= N; i++) {
            sum = prevSum + Long.parseLong(st.nextToken());
            prevSum = sum;
            if (sum % M == 0) {
                remainderZero++;
            } else {
                remainders[(int) (sum % M)]++;
            }
        }

        long answer = (remainderZero * (remainderZero + 1) / 2);
        for (int i = 0; i < M; i++) {
            answer += (remainders[i] * (remainders[i] - 1) / 2);
        }

        System.out.println(answer);
    }
}