[solved.ac 골드3] 10986_나머지 합 (파이썬)

https://github.com/mangbaam/CodingTest

GitHub - mangbaam/CodingTest: 프로그래머스, 백준 등 코딩테스트 풀이를 기록하는 저장소입니다.

밑의 사진을 클릭하면 문제 링크로 이동합니다

 

문제

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

예제

---

나의 풀이

N, M 의 범위가 상당히 크고 각 숫자의 최대 크기도 아주 크기 때문에 구간 마다 누적합을 구해서 전부 확인하는 것은 아주 아주 비효율적이다.

보통 이런 큰 숫자가 나오고 "나누어 떨어지는" 혹은 "n으로 나눈 나머지" 라는 조건이 붙으면 모듈러 연산(나머지 연산)의 특징을 잘 살리면 훨씬 빠르고 효율적인 코드를 작성할 수 있다.

 

예제로 주어진 1, 2, 3, 1, 2 를 가지고 누적합을 구해보면 1, 3, 6, 7, 9 가 된다. 각각을 3으로 나눈 나머지를 취해보면 1, 0, 0, 1, 0이 된다.

이것을 누적합을 구하면서 동시에 할 수도 있다.

1%3=1, 3%3=0, 3%3=0, 1%3=1, 3%3=0

위의 결과에서 보면 3으로 나누어 떨어져서 0이 된 부분이 3개 존재한다. 이는 각각 [0, 1] 구간합, [0, 2] 구간합, [0, 4] 구간합을 의미한다. 즉, 3으로 나누어 떨어지는 3개의 구간을 구할 수 있었다.

 

다음으로 생각해 볼 특징은 (A + B) % C 는 (( A % C) + (B % C)) % C 와 같다는 점이다. + 를 - 로 바꿔도 마찬가지 일 것이다.

(A - B) % C = (( A % C) - (B % C)) % C (이때 편의상 A > B 로 생각하자)

위 식에서 A 와 B 를 각각 특정 구간합으로 생각해보면 (( A % C) - (B % C)) % C 의 형태로 된 것이 우리가 구한 [1, 0, 0, 1, 0] 이 될 것이고, 같은 나머지 (예를 들어 0번 인덱스와 3번 인덱스 == 1)를 뺐을 때 0이 되기 때문에 0을 어떤 수로 나눠도 나머지는 0이다.

A = [1, 0, 0, 1, 0] 에서 A[3] - A[0] 은 0이고 이것은 [0, 3] 구간합 - [0, 0] 구간합을 의미한다. 즉 [1, 3] 구간합이 된다.

실제로 예제의 [1, 2, 3, 1, 2] 에서 [1, 3] 은 [2, 3, 1] 이 합이 6이 되면서 3으로 나눈 나머지는 0이 된다.

 

결국 위에서 구한 아예 0으로 나오는 3개의 구간 + 같은 숫자들의 조합(Combination) 수를 모두 더하면 경우의 수를 구할 수 있다.

c = lambda x: x * (x - 1) // 2

n, m = map(int, input().split())
li = list(map(int, input().split()))
last = 0
reminder = [0] * m
for i in li:
    last = (last + i) % m
    reminder[last] += 1

print(reminder[0] + sum([c(x) for x in reminder]))

그래서 실제 코드는 이렇게 아주 간단해진다. (위에서 설명한 것에서 조금 더 단축된 부분이 있긴 하지만)

 

last 는 누적합의 마지막 숫자를 의미하고, m 으로 나눈 나머지를 취한다.

나머지를 저장하는 리스트인 reminder 는 m 만큼의 크기를 가진다. m으로 나눈 나머지가 0 ~ m-1 이 나오기 때문이다.

reminder[last] 에 개수를 세어준다.

 

마지막에는 reminder[0] 에 reminder 각각의 숫자에 대해 Combination 의 합을 더한 값을 출력하면 정답이 된다.

pypy3 + sys.stdin.readline() 조합으로 시간을 더 줄일 수 있었다