문제
수 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으로 나누어 떨어지는 구간의 개수를 출력한다.
예제 입력 1
5 3
1 2 3 1 2
예제 출력 1
7
코드
import java.util.*;
import java.io.*;
public class Main {
public static void main(String[] args) {
Scanner sc = new Scanner(System.in);
int N = sc.nextInt();
int M = sc.nextInt();
long[] S = new long[N];
long[] C = new long[N];
long answer = 0;
S[0] = sc.nextInt();
for(int i=1; i<N; i++) {
S[i] = S[i-1] + sc.nextInt();
}
for(int i=0; i<N; i++) {
int remainder = (int) (S[i] % M);
if (remainder == 0) answer++;
C[remainder]++;
}
for(int i=0; i<M; i++) {
if (C[i] > 1) {
answer = answer + (C[i] * (C[i] - 1) / 2);
}
}
System.out.println(answer);
}
}
회고
무작정 코드로 들어가기 보다는 핵심아이디어를 짜고 슈도코드 작성하고 코드를 짜는게 훨씬 효율적이다.
핵심 아이디어는 두 가지로 나눌 수 있다.
- (A + B) % C는 ((A % C) + (B % C)) % C와 같다. 특정 구간 수들의 나머지 연산을 더해 나머지 연산을 한 값과 이 구간 합에 나머지 연산을 한 값은 동일하다.
- 구간 합 배열을 이용한 식 S[j] - S[i]는 원본 배열의 i + 1부터 j까지의 구간 합니다.
- S[j] % M의 값과 S[j] % M의 값이 같다면 (S[j] - S[i]) % M은 0이다. 즉, 구간 합 배열의 원소를 M으로 나눈 나머지로 업데이트하고 S[j]와 S[i]가 같은 (i, j)쌍을 찾으면 원본 배열에서 i+1부터 j까지의 구간 합이 M으로 나누어 떨어진다는 것을 알 수 있다.
Share article