백준 10986 <나머지 합> Python

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

10986번: 나머지 합

---

누적합을 사용하되 수학적 방법을 조금 사용해야 하는 문제입니다. mod 계산, 나머지 계산은 선형성을 띄기때문에 여타 계산 전에 나머지를 구하고 계산을 해도 여전히 나머지는 동일합니다. 즉 미리 누적합에 대해 나머지를 구한 뒤 나머지를 기준으로 리스트를 만들고 나머지를 사용해 콤비네이션 연산을 진행했습니다. 처음에는 itertools 의 combinations 를 구하고 리스트의 요소의 개수를 구하는 방법을 취했지만 시간초과 때문에 그냥 계산으로 처리했습니다.

---

import sys

from itertools import combinations

def solution(M, data_list):

    # 누적합

    subsum = 0

    # 구간합을 M 으로 나눈 나머지로 분류

    mod_list = [0 for _ in range(M)]

    for idx in range(len(data_list)):

        subsum += data_list[idx]

        mod_list[subsum%M] += 1

    # 카운트

    cnt = 0

    # 나머지가 0 인 경우는 우선 더함

    cnt += mod_list[0]

    # 차가 0 이 되는 경우를 조합해 콤비네이션 계산

    # -> mod_list 의 value 를 콤비네이션

    for idx in range(len(mod_list)):

        cnt += mod_list[idx] * (mod_list[idx]-1) // 2

    return cnt

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

data_list = list(map(int, sys.stdin.readline().split()))

print(solution(M, data_list))