정수 M으로 나누어 떨어지는 부분 구간합을 선형시간에 구하는 놀라운 방법

1. 문제

 

10986번: 나머지 합 (acmicpc.net)

10986번: 나머지 합

 

2. 풀이

 

누적합이야 prefix sum 방법으로 O(N)에 미리 구해놓고 [i,j]의 구간합은 O(1)에 구할 수 있는데

 

정수 M으로 나누어 떨어지는 구간합 [i,j]를 어떻게 하면 아주 빠르게 구할 수 있을까

 

가장 쉬운 방법은 모든 i,j에 대해 검사하는 이중 for문 방법이지만

 

N이 $10^{6}$이다 보니 $O(N^{2})$으로는 1초안에 통과할 수 있을리가 없다

 

그런데 [i,j]의 구간합을 어떻게 구하는지 생각을 해보면..

 

배열 A에서 0번 원소부터 i번 원소까지의 누적합 배열 prefix를 미리 O(N)에 구하고,

 

그러면 [i,j]의 구간합은 prefix[j] - prefix[i-1]로 O(1)에 구할 수 있다

 

그러니까 목표는 prefix[j] - prefix[i-1]이 정수 M으로 나누어 떨어지는 그러한 (i,j)의 개수를 구해야한다

 

-----------------------------------------------------------------------------------------------------------------

 

그런데 나머지를 구하는 모듈로 연산의 성질을 생각해보면,

 

https://deepdata.tistory.com/399

정수론 - 합과 곱은 왜 계속 나눠도 문제가 없는가? -

 

$$prefix[j] - prefix[i-1] (mod M) \equiv (prefix[j](mod M) - prefix[i-1](mod M)) (mod M)$$

 

prefix[j]와 prefix[i-1] 각각을 M으로 나눈 나머지끼리 뺀 값을 M으로 나눈 나머지가 

 

prefix[j]-prefix[i-1]을 M으로 나눈 나머지와 같다는 사실을 알고 있다.

 

이전에 O(N)으로 모든 prefix[i]를 이미 구해놨다.

 

그렇다면 모든 prefix[i]에 대해 M으로 나눈 나머지를 구한 다음에, 나머지가 서로 같은 값끼리 그룹을 지은다면 어떨까?

 

M으로 나눈 나머지는 0~M-1까지 가능하므로 크기 M인 배열에 각 그룹을 저장해두자.

 

그리고 각 그룹에서 2개를 뽑아 A,B라고 한다면,

 

A,B는 M으로 나눈 나머지가 서로 동일하다.

 

따라서 A와 B를 빼준 값은 M으로 나눈 나머지는 0이 된다.

 

그리고 A와 B를 빼준 값이 어떤 구간 [i,j]인지 특정할 수는 없지만, 어떤 구간 [i,j]의 구간합이 되는 것은 확실하다.

 

이 방법을 이용한다면 i,j를 모두 순회해보지 않더라도 M으로 나눈 나머지가 0인 구간합을 O(M)만에 구할 수 있게 된다.

 

import math
from sys import stdin

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

A = list(map(int,stdin.readline().split()))

prefix = [0]*(n+1)
group = [[] for _ in range(m)]
group[0].append(0)

for i in range(1,n+1):
    
    #prefix sum을 구해준다.
    prefix[i] = prefix[i-1] + A[i-1]
    
    #prefix sum 각각을 m으로 나눈 나머지를 구하고 각 그룹 내에 prefix sum값을 저장한다.
    group[prefix[i]%m].append(prefix[i])

answer = 0

#m으로 나눈 나머지가 동일한 그룹에서 2개를 뽑는 방법의 수를 전부 더하면,
#모든 경우에 [i,j]의 구간합이 m으로 나눈 나머지가 0인 경우가 된다.
for j in range(m):
    
    x = len(group[j])

    answer += math.comb(x,2)

print(answer)

 

또 한번 지렸다,.. 진짜 아이디어가 경이롭다