[누적합] 10986. 나머지의 합 (Python)

시도 1

deque로 풀 수 있을 것이라 생각했다..

[value, idx ] (value: 현재까지 값, idx : 현재 위치) 를 deque에 추가
하나씩 꺼내어 idx+1 번째의 값을 value에 더하고 m으로 나눈 나머지가 0인지 아닌지 체크

당연하게 시간초과..

시도 2 (정답)

list = [1,4], m으로 나눈 나머지 r 이라고 할 때, 다음과 같이 식을 작성할 수 있다

tmp1 = list[0] * m + r
tmp2 = list[1] * m + r

r의 값이 동일하다면,
tmp2 - tmp 1 = m * (list[1] - list[0])

 

이말은 즉,, 나머지가 같으면? m으로 나누어 떨어진다는 의미와 같다.
만약, list가 누적합 리스트라면, list[1] - list[0] 은 구간합을 의미한다.

 

따라서,,
나머지 같은 것 중에서 2개씩 뽑는 경우의 수와 같다 (단,, 나머지가 0일 때는 단독 가능하므로 추가로 더해줘야 함)

 

문제의 예시로 살펴보면
A = [1,2,3,1,2] 일 때, 누적합은 1, 3, 6, 7, 9 다
누적합의 나머지를 각각 계산하면? 1, 0, 0, 1, 0 이다.
나머지 list는 [3, 2, 0] ( 나머지 0은 3개, 1은 2개, 2는 0개)

nC2 = n*(n-1)//2 ( n 개 중 2개를 뽑는 경우의 수)이므로, 아래 식과 같이 각각에 대해 2개씩 뽑는 경우의 수를 계산하여 더한다.

for i in range(m):
    res+= r_list[i]*(r_list[i]-1)//2

++ 잊지말고 나머지가 0개일 경우, 단독 가능하므로 더해주기

 

전체 코드

import sys
input = sys.stdin.readline

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

r_list = [0]*m
prefix_sum = 0
for a in alist:
    prefix_sum += a
    r_list[prefix_sum%m]+=1

res = r_list[0]
for i in range(m):
    res+= r_list[i]*(r_list[i]-1) // 2
print(res)

 

 

문제링크

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