[백준 BOJ / Python] 10986번 나머지 합 (누적합)

문제

n개의 1차원 숫자 배열이 주어졌을 때, 연속된 부분합이 m으로 나누어 떨어지는 경우의 수를 구하는 문제이다.

풀이

• 부분합을 m으로 나누었을 때 나누어 떨어지는지가 문제임으로 나머지끼리의 연산을 통해 문제를 풀 수 있다.
• 특히 1차원 부분합은 누적합의 뺄셈을 통해 구할 수 있음으로
  $ partSum = prefixSum[j] - prefixSum[i]$  $( i < j ) $
  나머지가 같은 누적합의 뺄셈 계산을 통해 m으로 나눈 나머지가 0이 되는 부분합을 구할 수 있다.
  
  • 부분합을 구할 때는 항상 뒤에 위치한 j번째 누적합에 i번째 누적합을 빼야한다. (i < j) 
  • 즉, 나머지의 수가 같은 누적합의 조합을 통해 나머지가 0이 되는 부분합의 경우의 수를 구할 수 있다. $_nC_2 = \frac{n * ( n-1 )}{2}$
• 예제) 1차원 배열 [1,2,3,1,2], m = 3

• m으로 나누었을 때 나머지가 0이 될 경우의 수는 우선 나머지가 0으로 나온 3가지 경우의 수가 있다.
• 나머지가 0인 누적합은 3개이며 3,6,9의 조합을 통해 (6-3),(9-3),(9-6) 3가지 경우의 수를 추가할 수 있다.
• 나머지가 1인 누적합은 2개이며 1,7의 조합을 통해 (7-1) 1가지 경우의 수를 추가할 수 있다.
• 모든 경우의 수를 합하여 3 + 3 + 1 = 7 이 정답이 된다.

알고리즘

1. 나머지를 담을 m 길이의 리스트를 생성한다.
2. 누적합을 구하면서 해당 누적합을 m으로 나누었을 때 나온 나머지에 1을 더한다.
3. 나머지가 0이 되는 경우의 수를 ans에 담는다.
4. $_nC_2$ 조합 공식에 따라 누적합의 뺄셈을 통해 나머지 0을 만들 수 있는 경우의 수를 ans에 더한다.
5. ans를 출력한다.

Python 코드

import sys
input = sys.stdin.readline

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

r = [0] * m # m으로 나눈 나머지 담는 곳
prefixSum = 0
for i in range(n):
    prefixSum += a[i]
    r[prefixSum % m] += 1

ans = r[0] # 나머지가 0이 되는 경우의 수
for i in range(m):
    # nC2 = n(n-1)/2
    ans += r[i] * (r[i]-1) // 2
print(ans)

문제출처

10986번: 나머지 합