문제

아이디어

총 10^6개의 숫자가 주어지므로 시간복잡도는 O(n^2)을 넘을 수 없다.

누적합을 m으로 나눈 나머지 배열을 구하고, 이를 통해 정답을 구한다.

예제를 풀어보며 이해해보자.

5 3
1 2 3 1 2

위 그림에서 Mod 리스트는 다음을 의미한다:
Mod[x] = 나머지가 x인 수의 개수

또한 Num 리스트는 다음을 의미한다:
Num[A] = A번째 수 까지의 누적합을 m으로 나눴을 때의 나머지

A번째 수 X가 있을 때, 이 X를 마지막 요소로 하는 누적 합 중, m으로 나누어 떨어지는 누적 합의 개수는 Mod[Num[A]]와 일치한다.

다만,Num[A] == 0인 경우와 Num[A] != 0 인 경우의 누적 합을 확인하는 시점이 다름을 유의하자.

소스 코드

import java.io.*;
import java.util.*;

public class Main {
    static int n;
    static int m;
    static long[] mod;
    static long[] num;
    static long ans;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        n = Integer.parseInt(st.nextToken());
        m = Integer.parseInt(st.nextToken());
        ans = 0;

        num = new long[n];
        mod = new long[m];

        st = new StringTokenizer(br.readLine());
        num[0] = Long.parseLong(st.nextToken()) % m;
        mod[(int) num[0]]++;
        if(num[0] == 0) ans++;
        for (int i = 1; i < n; i++) {
            num[i] = (num[i - 1] + Integer.parseInt(st.nextToken())) % m;
            if (num[i] == 0) {
                mod[(int) num[i]]++;
                ans += mod[0];
            }
            else {
                ans += mod[(int) num[i]];
                mod[(int) num[i]]++;
            }
        }
        System.out.println(ans);
    }
}

num[i] == 0인 경우와 num[i] != 0인 경우의 mod와 ans 업데이트 시점이 다름을 확인할 수 있다.

채점 결과

정답을 정수형 자료형에 저장했는데, 최대 5*10^11의 크기를 가질 수 있으므로 Long으로 저장해주는 것이 옳았다.