백준 10986 : 나머지 합 (Python)

본문 링크

import sys
input=sys.stdin.readline

N,M=map(int,input().split())
L=list(map(int,input().split()))
answer=0
dp = [0] * (M)
dp[0]=1
total=0
for i in range(N):
    total+=L[i]
    dp[total%M]+=1
count=0
for i in dp:
    count+=i*(i-1)//2 #결국 M으로 나눠떨어지는 부분합을 구하는 것은 나머지가 동일한 idx 중 임의로 (nC2) 를 선택하는 것과 같다.
print(count)

"""
나머지가 같은 두 부분합을 고르면 두 구간은 결국 M의 배수가 된다.
나머지가 0인 경우는 부분합 자체가 M의 배수인 경우이므로 두 구간이 아니라 본인 구간도 될 수 있기 때문에
index가 0인 (부분합 0) 것을 포함
"""

📌 어떻게 접근할 것인가?

단순 구간합 문제라고 생각했다. 하지만 2차원 배열을 사용하니 메모리 초과. 2중반복문을 사용하니 시간초과가 나왔다.
따라서 우리는 더 간단한 아이디어를 통해 문제를 접근하고자 한다.

먼저 누적합 배열을 만들어 보자.

• 예제 입력
  5 3
  1 2 3 1 2

그럼 $SUM=[0,1,3,6,7,9]$ 가 된다.
이때 누적합의 배열을 $M$ 으로 나눈값을 저장해보면

$dp=[0,1,0,0,1,0]$ 가 되고
나머지가 0 과 1로 이루어진 배열을 개수로 나타내면
$dp=[4,2]$ 가 된다.
즉 나머지가 0 이 되는 경우 4개 , 나머지가 1이 되는 경우 2개
왜 이렇게 나누었는지는 누적합 배열을 보면된다.

나머지가 1이 되는 경우는 1과 7이다.
이때 7-1을 하면 6이된다.

즉 나머지가 같은 두 수를 빼면 M으로 나누었을때 나머지가 0이 된다.

따라서 $dp=[0]×M$ 만큼 배열을 만들어주고 dp[total%M]+=1 을 해줘서
나머지가 같은 누적합값을 dp 배열에 저장해준다.

이후 나머지가 같은 수들에서 2개를 임의로 구하기 때문에 $_nC_2$ 수식을 사용하였다.

✅ 코드에서 주의해야 할 점

• dp 배열의 크기는 M으로 설정한다.
• dp[0]=1로 잡아준다.