10986번 - 나머지 합

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

10986번: 나머지 합

문제를 보고 처음에 든 생각은 누적합을 이용해서 합 배열을 구한 후, 그 합 배열을 for 반복문으로 돌며 나머지 연산을 진행하여 나누어 떨어지는지 검사하는 것이다. 

# 나머지 합
import sys
input = sys.stdin.readline

N, M = map(int, input().split())
A = list(map(int, input().split()))

answer = 0
S = [0 for _ in range(N+1)]
S[1] = A[0]

for i in range(1, N):
    S[i+1] = S[i] + A[i]

for i in range(1, N+1):
    for j in range(i, N+1):
        if (S[j]-S[i-1]) % M == 0:
            answer += 1

print(answer)

합 배열 S를 선언하고 이중 for 문으로 i부터 j까지의 합을 구해서 나머지 연산을 하였다. 하지만 N의 범위가 10^6으로 적은 수가 아니기 때문에 당연하게도 시간 초과가 났다. 시간 초과를 해결하기 위해서 나머지 연산의 특성을 사용하였다. 이 내용은 Do it! 알고리즘 코딩테스트 with Python 을 참고하였다. 

1. (A + B) % C 와 ((A%C) + (B%C)) % C는 동일하다.
2. S[i] % M의 값과 S[j] % M의 값이 같다면, (S[j]-S[i]) % M은 0이다. 

문제를 푸는 데 2번 특성이 중요했다. 이 말이 무슨 뜻이냐면, 문제의 예제인 A가 1, 2, 3, 1, 2라고 할 때 합 배열 S는 1, 3, 6, 7, 9이다. 합 배열 S를 M의 값인 3으로 나머지 연산을 취하면 1, 0, 0, 1, 0이 된다. 이때 S[0]과 S[3]이 1로 값이 같다. 즉, (S[3]-S[0]) % 3은 0으로 A[1], A[2], A[3]의 합인 2+3+1이 3으로 나누어 떨어진다는 것을 의미한다. 

S[i]의 값이 0이라는 뜻은, 0부터 j까지의 구간의 합이 M으로 나누어떨어진다는 것을 의미하기 때문에 0의 개수를 세서 answer에 더해주고, 같은 값끼리 조합 연산을 진행하여 2가지를 뽑도록 하였다. 구간 i부터 j까지의 합이기 때문에 0~N-1 사이의 수 중에 2개를 뽑아야 한다.

# 나머지 합
import sys
input = sys.stdin.readline

N, M = map(int, input().split())
A = list(map(int, input().split()))

answer = 0
S = [0 for _ in range(N)]
cnt = [0 for _ in range(M)]

S[0] = A[0]

for i in range(1, N):
    S[i] = S[i-1] + A[i]

for i in range(N):
    remainder = S[i] % M
    if remainder == 0:
        answer += 1
    cnt[remainder] += 1

for i in range(M):
    # 조합
    if cnt[i] > 1:
        answer += (cnt[i] * (cnt[i]-1)) // 2

print(answer)

최종적으로 제출한 코드는 위와 같다.