[🥇3 / 백준 10986 / 파이썬] 나머지 합

 

10986번: 나머지 합

 

문제

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

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

입력

첫째 줄에 N과 M이 주어진다. (1 ≤ N ≤ 10⁶, 2 ≤ M ≤ 10³)

둘째 줄에 N개의 수 A₁, A₂, ..., A_N이 주어진다. (0 ≤ A_i ≤ 10⁹)

출력

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

 

예제 입력 5 3 1 2 3 1 2

예제 출력 7

 

풀이

 

모듈러 연산으로의 접근이 필요한 문제이다.

 

우선 구간의 합을 쉽게 구하기 위해 누적 합 배열(P)을 만든다.

예제의 경우 [1 2 3 1 2]이므로 [1 3 6 7 9] 가 될 것이다.

이제 누적 합 배열에서 인덱스 i, j (i<j)를 골랐을 때 구간의 합은 P[j] - P[i-1] 이 된다.

 

우리는 (P[j-1] - P[i]) % m = 0 이 되는 부분 구간들을 구해야 한다.

여기서 모듈러 연산의 경우 분배 법칙이 적용 가능하다.

따라서 P[j-1] % m = P[i] % m 인 것, 즉 누적합 배열에서 나머지가 같은 것들을 고르면 된다.

 

배열 P의 m으로 나눈 나머지들을 구하면 [1 0 0 1 0]이 된다.

나머지가 같은 것들을 2개 고르는 조합을 구하도록 한다.

0 : 3C2 = 3

1 : 2C2 = 1

단, 나머지를 담은 배열에서의 0들은 부분 구간이 아닌 첫 인덱스부터 더한 것이므로 따로 세주도록 하자.

그럼 3 + 1 + 3 = 7이 된다.

 

import sys, math
input = sys.stdin.readline

n,m = map(int,input().split())
arr = list(map(int,input().split()))

# 나머지만 담은 배열을 구한다. 예제의 경우 [3,2,0]
cnt_rem = [0]*m
cnt_rem[arr[0]%m]+=1

for i in range(1,n):
    arr[i] = (arr[i-1]+arr[i])%m
    cnt_rem[arr[i]]+=1

# 0의 개수 + 나머지가 같은 것들의 조합
ans = cnt_rem[0]
for i in cnt_rem:
    if i>=2:
        ans += math.comb(i,2)
print(ans)