10986. 나머지 합

스케치

• N이 주어졌을 때, 구간의 길이를 1부터 $N$까지 늘려가면서 모든 경우의 수를 조사하는 것은 시간 복잡도가 $2^N-1$ 로 매우 높음

초기 아이디어 (시간 초과)

• $O(N)$으로 구간 합을 구하고
• $O(n^2)$으로 각 구간 합값이 $m$ 으로 나누어떨어지는지 검사
  • 떨어지면 count++
• 근데 $n$ 이 $10^6$ 이라서 1초보다 오래 걸림 ($10^9$ = 1초)

import sys  
input = sys.stdin.readline  
  
n, m = map(int, input().split())  
a = list(map(int, input().split()))  
  
# 구간 합 구하기  
s = [0 for _ in range(n + 1)]  
for i in range(1, n + 1):  
    s[i] = (s[i-1] + a[i-1]) % m  
print(s)  
  
count = 0  
for i in range(1, n + 1):  
    for j in range(i, n + 1):  
        if (s[i-1] - s[j]) % m == 0:  
            count += 1  
  
print(count)

해결

1. 우선 이전과 같이 구간 합 S를 구한다.

2. 구간 합에 나머지 연산을 적용한 값 K도 구한다.

• $(a+b)\mathrm{mod}m=(a\mathrm{mod}m+b\mathrm{mod}m)\mathrm{mod}m$

3. S[i] - S[j]는 원본 리스트 A의 j+1부터 i까지의 구간 합이다.
4. S[i] % M과 S[j] % M의 값이 같으면 (S[i] - S[j]) % M = 0이다. 즉, S[i] % M = S[j] % M인 순서쌍 (i, j)이 있다는 것은 원본 리스트에서 j+1부터 i까지의 구간 합이 M으로 나누어떨어진다는 것을 알 수 있다.
5. $O(N+M)$의 복잡도로 최적화할 수 있다.

import sys  
input = sys.stdin.readline  
  
n, m    = map(int, input().split())  
a       = list(map(int, input().split()))  
s       = [0 for _ in range(n)]  
k       = [0 for _ in range(m)]  
count   = 0  
  
# 구간 합, 나머지값 구하기  
s[0] = a[0]  
k[s[0] % m] += 1  
for i in range(1, n):  
    s[i] = s[i-1] + a[i]  
    k[s[i] % m] += 1  
  
# 총 개수 구하기  
count += k[0]  
for i in range(m):  
    count += k[i] * (k[i] - 1) / 2  
  
print(int(count))