[GoldⅢ / BeakJoon / Java] 나머지 합

나머지 합 (백준 10986)

문제

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

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

입력

첫째 줄에 N과 M이 주어진다. (1 ≤ N ≤ 106, 2 ≤ M ≤ 103)

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

출력

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

풀이

N이 최대 1,000,000이기 때문에 O(N²)은 시간 초과가 난다. 그래서 누적합을 이용한 수학적 접근으로 해결했다.

구간 합(sumArr[j] - sumArr[i - 1])을 M으로 나눈 나머지가 0인 (i, j)의 값을 구해야한다. 수학적으로 정리하면 다음과 같다.

1. (sumArr[j] - sumArr[i - 1]) % M = 0
2. (sumArr[j] % M) - (sumArr[i - 1] % M) = 0
3. sumArr[j] % M = sumArr[i - 1] % M

이걸 토대로 구현해야할 코드를 단계적으로 생각해보자

1. 합 배열 sumArr를 생성할 때 M으로 나눈 나머지 결과를 저장한다

index	0	1	2	3	4	5
A[i]		1	2	3	1	2
sumArr[i]		1	3	6	7	9
sumArr[i] % M		1	0	0	1	0

2. sumArr[i] % M == 0 인 경우는 그 자체로 유효한 구간이므로 결과값에 더한다.

3. sumArr[j] % M == sumArr[i - 1] % M을 만족하는 (i, j)의 수를 구한다. 즉, 나머지 값이 같은 인덱스 2개를 뽑는 모든 경우의 수를 구하면 된다.

   나머지 값이 같은 인덱스의 수를 저장하기 위해remainArr 배열에 생성한다. 이때 M으로 나눈 나머지가 될 수 있는 수는 0 ~ M-1까지 가능하므로 크기가 M인 배열로 생성한다.
   remainArr 배열에 나머지 값이 같은 인덱스 중 2개를 뽑는 모든 경우의 수를 구한다. ₃C₂와 ₂C₂의 합을 구해 결과 값에 더해주면 된다.

   nCr : n개 중에서 r개를 뽑는 모든 경우의 수

소스 코드

public class Ma10986 {
    public static void main(String[] args) throws IOException {
        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[] sumArr = new long[N + 1];
        long[] remainArr = new long[M];
        long cnt = 0;

        st = new StringTokenizer(br.readLine());
        for (int i = 1; i < N + 1; i++) {
            sumArr[i] = Integer.parseInt(st.nextToken()) + sumArr[i - 1];
            int remain = (int) (sumArr[i] % M);

            if (remain == 0) cnt++;		 // 누적합이 M으로 나누어떨어질 때
            remainArr[remain]++;		 // 해당 나머지 카운트 증가
        }

		// 나머지가 같은 누적합 쌍 조합 수 구하기
        for (int i = 0; i < M; i++) {
            if (remainArr[i] > 1) cnt += (remainArr[i] * (remainArr[i] - 1) / 2);
        }

        System.out.println(cnt);
    }
}