백준 10986 나머지 합 파이썬

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

10986번: 나머지 합

N의 크기가 10**6이라 일반적으로 반복문을 돌려서 찾기에는 무리가 있어 누적합을 이용한 방법을 알아보다.

길이가 M개인 배열을 만든 뒤(DP)  입력 받은 리스트를 누적합을 해주고 M으로 나눈 나머지를 DP에 더해주면 DP[s[i]%M] +1 나머지의 목록이 나오는 것을 확인 할 수 있다.

나머지가 같은 누적합을 빼준다면 나머지가 0이 되기에 나머지가 같은 누적합 중에서 2개를 뽑으면 우리가 원하는 구간의 개수를 구할 수 있다. 그리고 나머지가 0인 부분은 그 자체로 만족하기에 더해준다.

조합은 math.comb를 사용했다.

예제 1번 ex)

누적합 : 1, 3, 6, 7, 9

나머지 : 1, 0, 0, 1, 0

7-1 = 6  (1~4)

9-3 = 6  (2~5)

9-6 = 3  (3~5)

3  (1~2)

6  (1~3)

9  (1~5)

import math
n,m = map(int,input().split())
s = list(map(int,input().split()))
dp = [0]*m
dp[s[0]%m]+=1
for i in range(1,n):
    s[i] = s[i-1]+s[i]
    dp[s[i]%m]+=1

re=dp[0]# 나머지가 0이니 더해준다. (0~i)
for i in range(m):
    re+=math.comb(dp[i],2) #2개를 뽑는 개수를 더해줌
print(re)