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

누적 합을 활용하는데, 최적화를 위한 테크닉이 필요합니다.

풀이

​문제 변환하기

먼저, 구간 $[1, i]$ $(1 \leq i \leq N)$에 대해, 구간 합을 $P_i$라고 정의합니다.

$P_i = \sum_{i=1}^N A[i]$

또, 구간 $[1, i]$에 대해 구간 합을 $M$으로 나눈 나머지를 $S_i$라고 정의합니다.

$S_i = P_i \mod M$

구간 $[i, j]$ $(1 \leq i \leq j \leq N)$에 대해, 구간 합은 $P_j - P_{i - 1}$로 정의됩니다.

문제에서 찾고자 하는 답은 "[i, j]의 구간 합을 M으로 나눈 나머지가 0이 되는 구간의 개수" 입니다. [i, j]의 구간 합을 M으로 나눈 나머지가 0이 되는 경우를 식으로 살펴봅시다.

$(P_j - P_{i-1}) \mod M \equiv (P_j \mod M - P_{i-1} \mod M) \mod M \equiv 0$

앞서 정의한 S를 사용해 다시 식을 나타내면, $(S_j - S_{i-1}) \mod M \equiv 0$이 됩니다. S는 M으로 나눈 나머지를 저장한 배열이므로, $(S_j - S_{i-1}) \mod M \equiv 0$이 성립하려면 $S_j = S_{i-1}$이어야 합니다.

따라서, 이 문제는 $S_j = S_{i-1}$이 되는 (i, j) 쌍의 개수를 세는 문제로 생각할 수 있습니다.

변환된 문제 풀기

$S_j = S_{i-1}$이 되는 (i, j) 쌍의 개수

앞서 [i, j]의 구간 합을 M으로 나눈 나머지가 0이 되는 구간의 개수를 찾는 문제를
$S_j = S_{i-1}$이 되는 (i, j) 쌍의 개수를 세는 문제로 바꿨습니다.

S는 $M$으로 나눈 나머지를 표현하므로, $0, \cdots, M - 1$의 값으로 표현됩니다. 따라서, $0, \cdots, M - 1$의 값이 S에 몇 번 등장하는지를 기록한 뒤, 각각의 나머지 $r$에 대해 $C(r, 2)$의 합을 구하면 됩니다.

구현

1) 나머지 개수 세기

freq_r = [0] * M
freq_r[0] = 1 # S_0은 항상 0이다. / S_0을 포함해 두어야, 시작점이 1인 모든 구간을 제대로 정답에 포함할 수 있다.
S = 0
for _ in range(N):
    S = (S + next(input) % M) % M
    freq_r[S] += 1

2) 답 구하기

answer = 0
for r in freq_r:
    answer += r * (r - 1) // 2
print(answer)

전체 코드

from sys import stdin
input = map(int, stdin.buffer.read().split()) # 입력 스트리밍
N = next(input)
M = next(input)
freq_r = [0] * M
freq_r[0] = 1 # S_0은 항상 0이다. / S_0을 포함해 두어야, 시작점이 1인 모든 구간을 제대로 정답에 포함할 수 있다.

S = 0
for i in range(N):
    S = (S + next(input) % M) % M
    freq_r[S] += 1

answer = 0
for r in freq_r:
    answer += r * (r - 1) // 2
print(answer)