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

문제

10986번: 나머지 합

수 N개 A1, A2, ..., AN이 주어진다. 이때, 연속된 부분 구간의 합이 M으로 나누어 떨어지는 구간의 개수를 구하는 프로그램을 작성하시오.

즉, Ai + ... + Aj (i ≤ j) 의 합이 M으로 나누어 떨어지는 (i, j) 쌍의 개수를 구해야 한다.

문제 해결 전략

기본 전략 설정

이 문제는 어떤 수열이 주어질때 특정 구간의 합이 M으로 나누어떨어지는 경우의 수를 찾는 문제이다.

누적합 테크닉을 이용해 매 구간 모든 요소를 더해 검사하는 비용을 절약할 수 있다.

문제의 특이사항 발견

특이사항은 수열의 길이가 최대 $10^6$, 수열을 이루는 요소들이 최대 $10^9$이 될 수 있다는 것이다.

전체 수열에 대해 값을 그대로 더해 누적합을 구할 경우 메모리가 펑 하고 터져버릴 것이기 때문에 값을 축소시켜 문제를 해결할 수 있는지 살펴보자.

구간 합이 M으로 나누어 떨어지는지 확인하기 위해서는 구간의 요소를 모두 더한 값을 M으로 나눈 나머지를 확인해야 한다. 이 때 모듈러 연산의 분배법칙에 의해 모든 수를 더한 후 M으로 나눈 나머지를 구하는 대신 각각의 요소를 M으로 나눈 값을 더한 후 나머지를 구할 수 있다.

(A + B) % p = ((A % p) + (B % p)) % p

따라서 수를 하나씩 더하며 누적합을 구할 때마다 M으로 나눠 누적합을 $0$ ~ $M - 1$ 사이의 값으로 유지하자. 누적합을 M으로 나눈 나머지를 확인하기에는 문제가 없다.

시간 절약하기

$(i, j)$ 구간의 합을 M으로 나눈 나머지가 0이 되려면 $(0, i - 1)$ 구간의 누적합과 $(0, j)$ 구간의 누적합을 각각 M으로 나눈 나머지가 서로 같아야 한다. 따라서 어떤 요소까지의 누적합이 주어졌을 때 이전 누적합 값 중 M으로 나눈 나머지가 같은 요소의 개수를 세면 M으로 나누어 떨어지는 구간의 개수를 구할 수 있다.

각각 요소까지의 누적합을 배열로 유지하고 매 요소마다 이전 요소들을 순회하며 나머지가 같은 요소의 수를 셀 수도 있겠지만 우리는 각각의 나머지 값이 몇번 등장했는지만 알면 되므로 이전 요소에서 $0$ ~ $M - 1$의 값이 각각 몇 번씩 등장했는지만 관리하면 누적합 배열을 순회하며 개수를 세는 시간을 절약할 수 있다.

정답 코드

import sys

def solution(n, m, a):
    r_cnt = [0 for _ in range(m)]
    r_cnt[0] = 1  # 0번 요소부터 시작하는 구간합의 경우를 포함하기 위해 나머지가 0인 경우를 1 추가한다.

    result = 0
    psum = 0
    for i in range(n):
        psum = (psum + a[i]) % m  # 누적합을 구할 때마다 m으로 나는 나머지를 저장한다.
        result += r_cnt[psum]  # 이전에 같은 나머지값이 몇번 등장했는지 확인해 추가한다.
        r_cnt[psum] += 1  # r_cnt 배열에서 현재 누적합 % m 값에 대한 카운트를 증가시킨다.

    return result

n, m = [int(i) for i in sys.stdin.readline().rstrip().rsplit()]
a = [int(i) for i in sys.stdin.readline().rstrip().rsplit()]

print(solution(n, m, a))

효율성 분석

매 요소까지의 누적합에 대해 M으로 나눈 나머지가 같은 요소가 몇 번 등장하는지 체크해 경우의 수를 구한다.

r_cnt 배열을 이용해 상수 시간에 누적합을 M으로 나눈 값이 몇 번 등장했는지 체크하거나 갱신할 수 있다.

위 작업을 N개 요소를 순회하며 수행하므로 전체 알고리즘은 선형 시간 복잡도를 띤다.