[코딩 테스트 연습(파이썬/Python)] 백준(BOJ) 10986번 _ 나머지 합

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

10986번: 나머지 합

[문제]

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

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

 

[입력]

첫째 줄에 N과 M이 주어진다. (1 ≤ N ≤ 106, 2 ≤ M ≤ 103)

둘째 줄에 N개의 수 A1, A2, ..., AN이 주어진다. (0 ≤ Ai ≤ 109)

[출력]

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

 

[아이디어]

1) index i에서의 누적 합을 구한다.(accumulate)

2) 구한 누적 합들의 m으로 나눈 나머지를 구한다.

3) 누적 합의 m으로 나눈 나머지가 0이라면 정답에 포함된다.

4) 누적 합이 m으로 나눈 나머지가 같은 것들 중 서로 다른 2개를 고를 때마다, 해당 두 위치의 차이만큼의 구간(사이 구간)에서 누적 합이 m으로 나누어 떨어진다.

 

from itertools import accumulate
import math
n,m=map(int,input().split())
a=list(map(int,input().split()))
a=accumulate(a)
d=[0]*m
for x in a:
    d[x%m]+=1
ans=d[0]
for k in d:
    ans+=math.comb(k,2)
print(ans)