백준 [10986] 나머지 합(Python3)

문제 링크 : https://www.acmicpc.net/problem/10986

10986번: 나머지 합

지문 해석

연속된 부분 구간의 합을 보고 바로 누적합을 써야 한다는 것을 알았다.

완전 탐색

누적합 배열이 있을 때, 한 누적합을 기준으로, 그 이전 위치의 누적합과 차를 비교해 M의 배수가 되는 쌍을 센다.
-> O(N^2)로 시간 초과

 

완전 탐색 최적화 - 탐색 범위 줄이기

누적합은 계속 커진다는 점을 이용해 두 누적합의 차가 M보다 작아지면 탐색을 중단한다.
하지만, M이 1이고, 모든 A의 값들이 0또는 1이라면 O(N^2)가 되어 시간초과

이 외에는 명확한 기준을 가지고 탐색 범위를 줄이는 것은 힘들다고 생각했기 때문에

새로운 누적합 배열을 만드는 방향으로 생각했다.

 

새로운 누적합 배열 생성 1

A에 대한 누적합이 아닌 새로운 종류의 누적합 배열을 생성한다.
A의 누적합이 M의 배수일 경우 카운트를 1하고, 이 값에 대한 누적합을 구한다.
이렇게 구해진 누적합을 가지고 할 수 있는게 없었기 때문에 다른 풀이로 넘어갔다.

새로운 누적합 배열 생성 2

여기서 놀라운 사실을 알 수 있었다. M으로 나눈 나머지가 같은 두 누적합을 빼주면 M으로 나누어 떨어진다는 것이었다. 이를 이용해, 카운팅 배열을 하나 만들어서 각 A의 누적합에 M 나머지 연산을 한 나머지 값을 1씩 증가시킨다.

두 구간의 합을 구하면 되기 때문에, a개 중 2개만 고르면 된다. 즉, aC2로 계산해 그 값들을 더해주면 된다.

총 시간 복잡도는 O(N + M)이다.

 

코드

N, M = map(int, input().split())
nums = list(map(int, input().split()))
accumulate = [0]
remain = [0 for i in range(M)]
remain[0] = 1
for num in nums:
    accumulate.append(accumulate[-1] + num)
    remain[accumulate[-1] % M] += 1
ans = 0
for r in remain:
    ans += (r * (r - 1)) // 2
print(ans)