[코딩인터뷰] 백준 10986번 나머지 합

백준 10986번 나머지 합

수학적인 추론을 이끌어 내야 풀 수 있는 문제입니다.

'A(i) ~ A(j) = S(j) - S(i-1)' 가 Prefix Sum 의 기본 공식이고 나머지가 0 이어야 한다는 점에 착안해서 새로운 공식을 도출해야 합니다.

'A(1) ~ A(i) = K (K = 정수)' 이고 'A(1) ~ A(j) = K' 라면, 'A(i) ~ A(j) = 0' 이어야 합니다. 왜냐하면 i 부터 j 까지 합으로 인한 변화가 없어야 같은 K 가 나올 수 있기 때문입니다.

그리고 'A(i) ~ A(j) = 0' 이라면 양변에 나머지 연산을 해도 동일해서 '(A(i) ~ A(j)) % M = 0' 입니다.

이점에 착안해서 Prefix Sum 의 기본 공식을 대입하면 'A(i) ~ A(j) = S(j) - S(i-1) = 0' 이 성립합니다. 이를 다시 정리하면 아래와 같이 됩니다.

'S(j) = S(i-1)' 을 다시 적으면 'S(j) % M = S(i-1) % M'

위 공식이 성립하는 경우의 수가 문제가 요구하는 정답이므로, 부분 합의 나머지를 저장해놓고 동일한 나머지가 나타난 경우를 정답에 추가하여 출력합니다.

import sys
readline = sys.stdin.readline
read = sys.stdin.read

N, M = map(int, readline().split())
A = list(map(int, readline().split()))
A.insert(0,0)
P = [0] * (N + 1)

# 부분합 나머지가 나타난 횟수 저장
S = [0] * M

answer = 0
for i in range(1, N+1):
    P[i] = (P[i-1] + A[i]) % M

    # 부분합 비교가 아닌
    # 1부터 i 까지 합의 나머지가 0 인 경우
    if P[i] == 0:
        answer += 1

    # 부분합 나머지 등장 횟수 증가
    S[P[i]] += 1

for j in range(M):
    # 부분합 나머지 등장 횟수 C 2 (조합 계산) = 등장 횟수 * (등장 횟수 - 1) / 2
    answer += (S[j] * (S[j] - 1)) // 2

print(answer)

코딩 테스트에서 이렇게 수학적인 생각을 하는건 어려울지도 모르겠네요...