문제는 생각보다 단순하게 풀리는 경우가 많다.

문제

수 N개 $A_1$, $A_2$, ..., $A_N$이 주어진다. 이때, 연속된 부분 구간의 합이 $M$으로 나누어 떨어지는 구간의 개수를 구하는 파이썬 프로그램을 작성하라. 즉, $A_i + ... + A_j (i \leq j)$ 의 합이 M으로 나누어 떨어지는 (i, j) 쌍의 개수를 구해야 한다.

입력

첫째 줄에 N과 M이 주어진다. (1 ≤ N ≤ 106, 2 ≤ M ≤ 103) 둘째 줄에 N개의 수 $A_1$, $A_2$, ..., $A_N$이 주어진다. (0 ≤ $A_i$ ≤ 109)

출력

첫째 줄에 연속된 부분 구간의 합이 $M$으로 나누어 떨어지는 구간의 개수를 출력한다.

문제 해결 과정

1. 문제를 잘못 이해함 처음에는 "연속된" 이라는 것이 단순히 부분 누적합을 의미하는 줄 모르고 수의 대소관계 상 크고 작음을 이야기하는 줄알고 별 생각 없이 진행했다가 삽질을 조금 했다. 입/출력 예시를 제대로 읽어보고 직접 손으로 계산해 문제를 제대로 이해한건지 검토해보자.

2. 초기 아이디어 당연히 문제 인식이 잘못됐으므로 누적합 리스트를 만들고, 누적합 리스트를 순회하면서 대소관계가 맞는 부분 누적합에 대해 M으로 나는 나머지가 0인 경우 count를 증가하게 하려고 했다. 하지만 무슨 짓을 해도 입출력 예시가 만들어 지지 않는 것을 뒤늦게 깨닿고는 수정하게 되었다.

   누적합 계산 이후 부분 누적합에 대해 작업변수 i, j 두개의 for 문을 이용해 cumsum_list[j] - cumsum_list[i-1]을 $M$으로 나눈 나머지가 0인 경우 count를 증가하게 했다. 하지만 백준에서 채점해 보았을 때 시간 초과 문제에 직면하게 되었다.

3. 시간복잡도 줄이기 챗지피티를 이용해 시간복잡도를 줄일 수 있는 방법에 대해 물어봤다. 그러니 여러개를 알려주는데, 그 중 누적합 계산 직후 remainer_count 리스트에서 해당 누적합에 대한 나머지가 발견된 횟수를 count에 더해준다. 그 후, remainer가 0이면 count를 증가시키고 그것과 무관하게 나머지에 해당하는 remainer_count의 값을 1 증가시킨다. 이 로직의 핵심은 remainer_count를 이용해 중복을 피한 "현재 누적합"의 나머지를 체크하는 것이다.

오늘의 교훈

급할수록 차근차근

import sys

input = sys.stdin.readline
N, M = map(int, input().split())
a = list(map(int, input().split()))

count = 0
cumsum = 0
remainder_count = [0] * M

for num in a:
    cumsum += num
    remainder = cumsum % M
    count += remainder_count[remainder]
    if remainder == 0:
        count += 1
    remainder_count[remainder] += 1

print(count)