[백준/Python] 10986 - 나머지 합

문제

수 N개 A1, A2, ..., AN이 주어진다. 이때, 연속된 부분 구간의 합이 M으로 나누어 떨어지는 구간의 개수를 구하는 프로그램을 작성하시오.

즉, Ai + ... + Aj (i ≤ j) 의 합이 M으로 나누어 떨어지는 (i, j) 쌍의 개수를 구해야 한다.

입력

첫째 줄에 N과 M이 주어진다. ($1 ≤ N ≤ 10^6$, $2 ≤ M ≤ 10^3$)

둘째 줄에 N개의 수 A1, A2, ..., AN이 주어진다. ($0 ≤ Ai ≤ 10^9$)

출력

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

풀이
누적합을 구해준 후 구간합을 구하는 원리를 사용하면 되지만 실제로 모든 조합의 구간합을 구하면 시간초과가 발생한다
구간 합을 m으로 나누었을 때 나머지가 0인 경우는 2가지 경우가 존재한다

1. 구간합을 구성하는 두 요소가 모두 m의 배수, 즉 나머지가 0인 경우
2. 구간합을 구성하는 두 요소를 m으로 나눈 나머지가 0이 아니지만 두 요소의 나머지가 동일한 경우, 즉 두 요소를 뺄셈하면 3의 배수가 되는 경우

그리고 1번의 경우에는 누적합 값 자체가 m의 배수가 되는 경우와 두 요소가 모두 m의 배수가 되는 경우가 섞여있다
예를 들어

m = 3
  1 2 3 1 2
=>1 3 6 7 9

위와 같을 때 3,6,9는 다른 요소와의 뺄셈 없이 그 자체로 3의 배수가 된다
이 경우는 누적합 값 자체가 m의 배수가 되는 경우는 누적합 리스트를 만드는 과정에서 계산해줬다

그 외에 누적합끼리 값을 뺄셈하는 것이 필요한 경우를 처리하기 위해서 나머지 값을 인덱스로하는 m개의 원소를 갖는 리스트를 선언해줬다

나머지가 0이면 인덱스 0에 1을 더해주고, 나머지가 1이면 인덱스 1에 1을 더해주는 방식이다
이후 나머지 값이 같은 것들끼리 경우의 수를 계산하고 전체를 합산하면 된다
m이 3일 때 나머지가 0인 값들이 3개가 있고 1인 값들이 2개, 2인 값들이 3개 있다면
위에서 누적합 값 자체가 m의 배수인 경우의 수 + 3C2+2C2+3C2를 계산 한 값이 최종 결과값이다

코드

import sys

n, m = map(int, sys.stdin.readline().split())
a = list(map(int, sys.stdin.readline().split()))

prefix_sum = []
answer = 0
num = 0
etcs = [0]*m
for elem in a:
    num += elem
    prefix_sum.append(num)
    if num%m == 0:
        answer += 1
    etcs[num%m] += 1

for elem in etcs:
    answer += elem*(elem-1)//2
print(answer)