[백준/10986] 나머지 합 (Python)

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

10986번: 나머지 합

 

Difficulty : Gold 3

 

Status : Solved

 

Time : 00:16:21

 

---

 

문제 설명

 

 

수 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으로 나누어 떨어지는 구간의 개수를 출력한다.

 

입력 예시

 

5 3
1 2 3 1 2

 

출력 예시

 

7

 

 

---

 

풀이

 

N이 최대 10^6이므로, 전부 비교하는 풀이(O(N^2))는 시간 초과를 야기한다. 따라서 누적 합을 이용한 풀이로 귀결된다.

 

내가 처음 시도한 풀이, 그리고 다음에 소개할 풀이 둘 다 핵심은 동일하다. i까지의 누적합을 S[i]로 두면, i부터 j까지의 구간은 S[j] - S[i-1]이 된다. 나머지 연산의 경우 분리법칙이 성립하므로, ( S[j] - S[i-1] ) % M = S[j] % M - S[i-1] % M이 된다.

 

내가 시도한 풀이는 O(N^2)을 최대한 근사하고자 하는 풀이이다. 전체 누적합에서 앞구간을 1씩 줄여나가며, 나머지 리스트가 0을 카운트하는 경우를 전부 저장하였다. 이는 나머지 리스트를 회전시킴으로써 구현할 수 있다. 이 경우 시간복잡도는 대략 O(NM)이 된다.

 

풀이 코드 1

from collections import deque
import sys
input = sys.stdin.readline

N, M = map(int, input().split())
cvt = lambda x : int(x) % M
num_list = list(map(cvt, input().split()))

mod_list = deque([0]*M)
sum_val = 0
for i in range(N) :
  sum_val = (sum_val + num_list[i]) % M
  mod_list[sum_val] += 1

answer = 0
for i in range(N) :
  answer += mod_list[0]
  mod_list[num_list[i]] -= 1
  mod_list.rotate(-num_list[i])

print(answer)

 

또한 좀 더 정석적으로 접근하자면, 두 누적합의 나머지가 동일한 경우라면 그 구간합의 나머지는 0이 될 수 있다는 의미이다. 따라서 누적합의 나머지를 전부 카운트하여 저장하고, 그 경우의 수를 구하면 된다. 위와 같은 나머지를 가진 구간 N 중 2개를 순서 상관 없이 뽑는 경우의 수는 C(N, 2) = N*(N-1) // 2임을 이용하여 풀 수 있겠다. 이 경우는 O(N + M)의 시간복잡도를 가지게 된다.

 

풀이 코드 2

N, M = map(int, input().split())
num_list = list(map(int, input().split()))
mod_list = [0]*M
sum_val = 0
for i in range(N) :
  sum_val = (sum_val + num_list[i]) % M
  mod_list[sum_val] += 1
answer = mod_list[0]
for m in mod_list :
  answer += m * (m-1) // 2

print(answer)

풀이 완료!