10986번 나머지 합

https://www.acmicpc.net/problem/10986

시간초과 코드

n의 범위만 더 작았더라면 이것도 백퍼 맞았을 것 같다.

n, m = map(int, input().split())
num = list(map(int, input().split()))
prefix_sum = [0]*(n+1)
sum = 0
for i in range(1, n+1):
    sum += num[i-1]
    prefix_sum[i] = sum

cnt = 0
for i in range(1, n+1):
    for j in range(i):
        if (prefix_sum[i] - prefix_sum[j]) % m == 0:
            cnt += 1
print(cnt)

아이디어

모든 구간들의 누적합을 구해서 나눠도 되지만(위 코드처럼) n의 범위가 1000000으로 매우 커서 시간 초과가 난다.

1. 모든 구간들의 누적합을 m으로 나누어서 나온 나머지 인덱스에 넣는다.
   ex) 누적합 : 3 5 6 20, m=3일 때
   3으로 나누어떨어지는 것(나머지=0)은 2개
   3으로 나누었을 때 나머지가 1인 것은 0개
   3으로 나누었을 때 나머지가 2인 것은 2개
   즉, remainder = [2, 0, 2]가 된다.
   => 이렇게 한 이유는
   (나머지가 2인 구간) - (나머지가 2인 구간) = (나머지 0)
2. remainder 같은 인덱스에서 숫자 두 개씩 뽑으면 무조건 나누어떨어지기 때문에 모든 인덱스에서 nC2를 구해주는 것이 총 경우의 수가 된다.
   nC2 = n*(n-1) // 2

코드

n, m = map(int, input().split())
num = list(map(int, input().split()))
remainder = [0]*m
sum = 0

for i in range(n):
    sum += num[i]
    remainder[sum%m] += 1

answer = remainder[0] # 이미 나머지가 0이면 무조건 나누어 떨어짐
# 나머지가 1인 구간 두 곳 뽑아서 빼면 나누어 떨어짐
# (나머지가 1인 구간) - (나머지가 1인 구간) = (나머지 0)
for i in remainder:
    answer += i*(i-1) // 2
print(answer)