[Java+Python]10986번, 나머지 합

목차

 

누적 합

 

[Java+Python]11659번, 구간 합 구하기 4

[Java+Python]2559번, 수열

[Java+Python]16139번, 인간-컴퓨터 상호작용

[Java+Python]11660번, 구간 합 구하기 5

[Java+Python]25682번, 체스판 다시 칠하기 2

 

문제

 

수 $N$개 $A_1, A_2, ..., A_N$이 주어진다.

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

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

 

입력

 

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

둘째 줄에 $N$개의 수 $A_1, A_2, ..., A_N$이 주어진다. $(0 ≤ A_i ≤ 10^9)$

 

출력

 

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

 

풀이

 

생각보다 까다로운 문제라 시도하고 푸는데 시간이 좀 걸렸다.

 

그 과정도 중요하긴 하지만 결과만 적자면 다음과 같은 로직으로 풀면 된다.

 

1. 입력으로 주어진 n개의 숫자를 받아 누적합을 저장할, 길이 n의 배열 sum을 만든다.
2. 누적합 중 0 ~ (m - 1)의 나머지를 갖는 요소의 개수를 저장할, 길이 m의 배열 cnt를 만든다.
3. 각 배열을 채운 뒤 결괏값을 나머지가 0인 요소의 개수로 초기화해 준다.
4. 나머지가 같은 요소(누적합)의 차를 구하면 해당 결과는 m으로 나누어 떨어지게 된다.
5. 따라서 나머지가 같은 요소 중 두 개를 순서에 상관없이 뽑는 경우의 수를 구해 cnt 배열에 더해준다.
   1. 여기서 순서에 상관없이 뽑는 이유는 두 수를 골라 큰 값에서 작은 값을 빼기만 하면 되기 때문이다.

여기서 4번에 다다르기가 쉽지 않았던 것 같다.

 

하지만 늘 그렇듯 이렇게 정리하고 보면 불과 5-6줄에 불과한 로직이니 조금 부끄럽기도 하네.

 

계속해서 자바와 파이썬으로 구현해 보자.

 

Java

 

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

public class Prob10986 {

	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[] sum = new long[n];
		long[] cnt = new long[m];

		st = new StringTokenizer(br.readLine());
		for (int i = 0; i < n; i++) {
			int num = Integer.parseInt(st.nextToken());
			if (i > 0) {
				sum[i] = sum[i - 1];
			}
			sum[i] += num;

			int remain = (int)(sum[i] % m);
			cnt[remain]++; // 나머지 값의 개수에 따라 누적합 분류
		}

		long result = cnt[0];  // 누적합 중 나머지가 0인 요소의 개수로 초기화

		for (int i = 0; i < m; i++) {
			result += (cnt[i] * (cnt[i] - 1)) / 2;  // 나머지가 같은 경우 그 차이는 m으로 나누어 떨어진다.
		}

		System.out.println(result);
	}
}

 

Python

 

import sys

n, m = map(int, sys.stdin.readline().rstrip().split())
lst = list(map(int, sys.stdin.readline().rstrip().split()))
cnt = [0] * m

for i in range(n):
    if i > 0:
        lst[i] = lst[i - 1] + lst[i]
    remain = lst[i] % m
    cnt[remain] += 1

result = cnt[0]

for i in range(m):
    result += (cnt[i] * (cnt[i] - 1)) // 2

print(result)

 

Performance