문제

[백준] 나머지 합 (Gold 3)

요약

N개의 수가 주어지면, 연속된 부분 구간의 합이 M으로 나누어 떨어지는 경우의 개수를 구하라. (1 ≤ N ≤ 1,000,000, 2 ≤ M ≤ 1,000)

분류

  • 누적합
  • 수학 (나머지 정리)

풀이

내 풀이

N의 크기가 크기 때문에 O(N) 시간에 끝내야 한다. 중등 수학 과정인 나머지 정리를 활용하면 쉽게 해결할 수 있다.

아이디어는 다음과 같다.

  • prefix_sum[i] = arr[0]부터 arr[i]까지의 합
  • mods[i] = prefix_sum[i] % m
  • [0, a] 구간의 합의 나머지와 [b, n - 1]까지의 합의 나머지가 같다면 구간 [a, b] 의 합의 나머지는 나머지 정리에 의해 0이 된다.
  • 따라서 나머지가 같은 부분 수열 두 개를 고르면 된다.
n, m = map(int, input().split())
arr = list(map(int, input().split()))

acc = 0
prefix = [0] * n
mods = [0] * m
ans = 0

for i in range(n):
    acc += arr[i]
    prefix[i] = acc
    mods[acc % m] += 1
    if acc % m == 0:
        ans += 1


for i in range(m):
    if mods[i] < 2:
        continue
    ans += mods[i] * (mods[i] - 1) // 2 # 나머지가 같은 구간 2개를 고르는 경우의 수 : mods[i] COMBINATION 2
print(ans)