[백준 10986 파이썬] 나머지 합 (골드 3, 누적합)

알고리즘 유형 : 누적합
풀이 참고 없이 스스로 풀었나요? : X

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

소스 코드(파이썬)

import sys
input = sys.stdin.readline

N, M = map(int, input().split())
arr = list(map(int, input().split()))
p_sum = [0]*(N+1) # 부분합 리스트 (단, M으로 나눈 나머지로 저장)

# count[idx]는 1부터 x까지의 부분합이 담긴
# (정확하는 M으로 나눈 나머지가 담긴)
# p_sum의 원소들 중에, 그 값이 idx인 것의 개수를 의미함
count = [0]*(M+1)

# p_sum과 count 구하기
for i in range(N):
    p_sum[i+1] = (p_sum[i] + arr[i]) % M
    count[p_sum[i+1]] += 1

# 1부터 x까지의 부분합을 M으로 나눈 나머지가 0이라면,
# 그 부분합 자체로 하나를 카운팅해줘야됨
ans = count[0]

# 그 외의 부분합들을 처리하는 구문임.
# 만약 어떤 두 부분합(1부터 x까지의)을 M으로 나눈 나머지가 같고
# 각각 1부터 a, b 까지의 부분합이라고 가정하면,
# a+1부터 b까지의 부분합은 M으로 나누어떨어진다.
# 이 부분은 수학적인 논리이므로 잘 생각해보자.
for i in range(M+1):
    # 뽑는 개수가 2인 combination의 변형식임
    ans += (count[i] * (count[i]-1)) // 2
    
print(ans)

풀이 요약

1. 약간의 수학적인 마인드가 필요한 문제이다.

   우선 이 문제는 단순히 부분합을 구한 뒤 가능한 모든 경우의 부분합을 구해서 각각을 M으로 나눠 나머지를 검사하는 것은 O(N^2)인데 이 문제에서 N은 10^6까지가 범위이므로 TLE다.

   그래서 맨 처음 부분합을 구할 때 살짝 손을 봐서 O(N)으로 풀려고 생각을 해봐야한다.

2. 	0	1	2	3	4	5
   arr	0	1	2	3	1	2
   p_sum	0	1	3	6	7	9
   p_sum%M	0	1	0	0	1	0

우선 부분합을 평범하게 구한 뒤에, 그 원소들을 싹다 M으로 나눈 나머지로 처리해보자.

이제 여기서 수학적 지식이 활용되는데, 1부터 a까지의 부분합과 1부터 b까지의 부분합의 M으로 나눈 나머지(각각 a_sum%M[a], a_sum%M[b])가 서로 같다면, a+1부터 b까지의 부분합은 M으로 나누어 떨어진다

즉, 우리는 p_sum%M 리스트에서 값이 같은 것들을 찾아서, 그 중에서 2개를 뽑는 경우의 수를 구하면 된다. 단, 그 값이 0인 경우는 그냥 그 부분합(1부터 idx까지의) 하나만으로도 조건에 부합하는 카운팅이 되므로 빼먹지말자.

3. count 리스트를 만든다. count[idx]는 p_sum%M에서 값이 idx인 것의 개수를 의미한다.

   count 리스트를 만들고, 이 것을 순회하면서 그 값 x에 대해 $_{x}C_{2}$를 ans에 더해줘나가면 된다.

   $_{x}C_{2}$는 x(x-1) / 2 이다. 원래 공식 x!/r!(x-r)! 을 약분해보면 저 식이 나온다.

배운 점, 어려웠던 점

• 전에 배웠던 누적합 테크닉에서 벗어나지 않고, 약간의 수학적 지식을 접목하면 풀 수 있는 문제인데 너무 꼬아서 생각해서 접근하려는 습관이 남아있는 것 같다.

  저번에 배운 적 있는 고난이도 문제 접근 방법을 잊지 않고 적용하려 노력해야겠다.