문제 확인
10986번: 나머지 합
수 N개 A1, A2, ..., AN이 주어진다. 이때, 연속된 부분 구간의 합이 M으로 나누어 떨어지는 구간의 개수를 구하는 프로그램을 작성하시오. 즉, Ai + ... + Aj (i ≤ j) 의 합이 M으로 나누어 떨어지는 (i, j)
www.acmicpc.net
- 시간 제한 : 1초
- 메모리 제한 : 256MB
나의 풀이
- 부분 합 개념을 이용해서 풀었다.
- 이때, 부분 합을 계산한 배열을 $i, j$로 순차적으로 접근할 경우 $O(N^2)$으로 시간 초과가 발생한다. 따라서 다른 방식으로 풀이해야 하고, 풀이는 다음과 같다.
미리 구해둔 prefix sum에서 부분 합을 구하려면 P[j] - P[i - 1]의 과정을 거쳐야 한다.
이때, P[j]와 P[i-1]이 각각 M으로 나눴을 때의 나머지가 동일하다면 빼기 연산 후 나머지는 0이 돼 M으로 나눠떨어지게 된다.
결국 M으로 나눠떨어지는 부분합을 구하는 것은 나머지가 동일한 idx 중 임의로 $\binom{n}{2}$를 선택하는 것과 같다.
- [1] [2] [3] [1] [2] -> [0] [1] [3] [6] [7] [9] -> [0] [1] [0] [0] [1] [0] -> [0 : 4, 1 : 2] -> 6 + 1 -> 7
import sys
n, m = map(int, input().split())
num_list = list(map(int, sys.stdin.readline().rstrip().split()))
remainder_info = [0 for _ in range(m)]
remainder_info[0] = 1
total = 0
for i in range(n):
total += num_list[i]
r = total % m
# 나머지 값에 따라서 idx 정보 저장
remainder_info[r] += 1
count = 0
for i in remainder_info:
count += i*(i - 1) // 2
print(count)수정
2022-10-05 : 디버깅 과정에서 사용한 prefix_sum 제거