[백준/파이썬] 10986번: 나머지 합

문제

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

풀이

• 연속된 부분 구간 합이 M으로 나누어 떨어지는 구간 개수

문제를 보자마자 바로 누적합 알고리즘을 사용해 다음과 같이 코드를 작성했습니다.

import sys
input = sys.stdin.readline

n,m = map(int,input().split())
arr = list(map(int,input().split()))

tmp = 0
pre = [0]
for i in arr:
    tmp += i
    pre.append(tmp)

res = []
for i in range(1, n+1):
    for j in range(i):
        res.append(pre[i]-pre[j])

cnt = 0
for idx in res:
    if idx % m == 0:
        cnt += 1
print(cnt)

메모리 초과 ,,

정답을 모두 res에 저장해서 메모리 초과가 발생했나 싶어 res 배열에 저장하는 것이 아닌 다른 방법을 사용했습니다.
바로 누적합 값이 m으로 나누어 떨어지면 cnt 값을 증가하게 코드를 수정했지만 시간 초과 😢

다시 처음부터 생각해봅시다..

구간합이 M으로 나누어 떨어지려면, 두 누적합의 나머지가 같아야 합니다.
즉, 구간 [i,j]의 합이 pre[j] - pre[i-1]이고 이것이 M으로 나누어 떨어져야 합니다.
$pre[j] (mod M) = pre[i-1] (modM)$

누적합의 나머지를 M으로 나누었을 때 동일한 값이 등장하면 해당 구간의 합이 M으로 나누어 떨어진다.

배열을 순회하며 각 원소를 더해 누적합을 계산하고 이를 M으로 바로 나누어 나머지를 구합니다.

tmp = 0
for i in arr:
	tmp += i # 누적합
	mod = tmp % m # 누적합이 m으로 나눈 나머지

나머지가 mod인 누적합이 몇 번 등장했는지 기록

pre = [0] * m # 나머지가 출현한 횟수
pre[0] = 1 	# 누적합이 처으부터 m으로 나누어 떨어지는 경우를 포함하기 위해
			# 나머지가 0인 경우 +1
res = 0

for i in arr:
	tmp += i
    mod = tmp % m
    
    res += pre[mod] # 현재 나머지와 동일한 이전 카운트만큼 결과에 추가
    pre[mod] += 1 # 현재 나머지 등장 횟수 증가

코드

import sys
input = sys.stdin.readline

n,m = map(int,input().split())
arr = list(map(int,input().split()))

tmp = 0
pre = [0] * m
pre[0] = 1
res = 0

for i in arr:
    tmp += i
    mod = tmp % m

    res += pre[mod]
    pre[mod] += 1

print(res)