나머지 합 - 10986

[Gold III] 나머지 합 - 10986

문제 링크

분류 1️⃣

수학(math), 누적 합(prefix_sum)

문제 설명 2️⃣

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

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

입력 3️⃣

첫째 줄에 N과 M이 주어진다. (1 ≤ N ≤ 10⁶, 2 ≤ M ≤ 10³)

둘째 줄에 N개의 수 A₁, A₂, ..., A_N이 주어진다. (0 ≤ A_i ≤ 10⁹)

출력 4️⃣

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

예제 입력 1 🔽

5 3
1 2 3 1 2

예제 출력 1 🔽

7

---

풀이 🔓

1. 구간의 합이 M으로 나누어 떨어지는 경우의 수를 찾는 다는 말은 (s[i] - s[j-1]) % m == 0인 경우의 수를 찾으라는 말과 같다. 이때 s는 구간 합 배열,

2. (s[i] - s[j-1]) % m == (s[i] % m - s[j-1] % m) -> s[i] % m == s[j-1] % m

3. 즉, 합 배열에서 나머지 연산을 취한 새로운 나머지 배열에서 나머지가 같은 요소 중 두 개를 뽑아야 함 + 합 배열에서 나누어 떨어지는 경우의 수 (0, i) % m == 0

4. N개의 수에서 2개를 뽑는 경우의 수는 n*(n-1) // 2

코드 📃

import sys
from itertools import combinations

input = sys.stdin.readline
n, m = map(int, input().split())
a = list(map(int, input().split()))

s = [0 for _ in range(n)]
s[0] = a[0]
for i in range(1, n):
  s[i] = s[i-1] + a[i]


cnt = 0 # (i, j) 구간이 m으로 나누어 떨어 지는 횟수
s_m = [0 for _ in range(m)]
for i in range(n):
  r = s[i] % m
  if r == 0: # (0 ~ i)까지의 구간
    cnt += 1 
  s_m[r] += 1

for i in s_m:
  if i == 0:
    continue

  cnt += i*(i-1)/2
  # i개 중 2개를 뽑는 경우의 수 : i*(i-1)/2

print(int(cnt))