[python] baekjoon 백준 10986 나머지합

문제

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

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

 

입력

첫째 줄에 N과 M이 주어진다. (1 ≤ N ≤ 106, 2 ≤ M ≤ 103)

둘째 줄에 N개의 수 A1, A2, ..., AN이 주어진다. (0 ≤ Ai ≤ 109)

 

출력

첫째 줄에 연속된 부분 구간의 합이 M으로 나누어 떨어지는 구간의 개수를 출력한다.

 

풀이

부분 합을 계속 구했더니 시간 초과가 발생하였고, 결국 ChatGPT 도움을 받아서 해결했다.

그런데 봐도 봐도 이해가 잘 안돼서.. 결국 이해 했지만, 나와 같은 사람이 있다면 빠르게 도움을 줄 수 있도록 설명한다.

 

def count_subarrays_with_sum_divisible_by_M(arr, N, M):
    remainder_counts = [0] * M
    prefix_sum = 0
    result = 0

    for i in range(N):
        prefix_sum = (prefix_sum + arr[i]) % M
        if prefix_sum == 0:
            result += 1
        result += remainder_counts[prefix_sum]
        remainder_counts[prefix_sum] += 1

    return result

# 입력
N, M = map(int, input().split())
arr = list(map(int, input().split()))

# 함수 호출 및 출력
print(count_subarrays_with_sum_divisible_by_M(arr, N, M))

 

O(N)으로 풀리며 N번 루프를 반복한다.

루프는 

1. arr[i] 원소를 꺼내어 M으로 나눈 나머지를 구하고(prefix_sum)

2. 이때 prefix_sum이 0이라면 res에 1을 더한다.(나누어 떨어지기 때문에)

3. 다음이 중요한데 remainder_counts에 prefix_sum의 값을 인덱스로 가진 만틈 res에 더한다.

4. 그리고 remainder_counts에 prefix_sum 인덱스를 1 더한다.

 

3번이 왜 그럴까 계속 들여다 보고 있었는데 부분합이라는 사실 덕분에 알 수 있었다.

예시처럼 1 2 3 1 2 라는 배열이 있고 3번째(i=2) 루프를 돌때

remainder = [1, 1, 0], prefix_sum = 0 이다.

이때는 prefix_sum = 0 이기 때문에 res += 1 이 되면서

이전에서 루프에서 1 + 2 = 3 이라는 값이 존재했기 때문에 remainder[0] = 1 이므로 res += 1 이 된다.

1 2 그리고 1 2 3 그리고 3 루프 3개가 있기 때문에 결국 res 값은 3이 된다.

 

remainer_counts 라는 변수의 명칭법과 더불어 누적합에 대한 이해가 조금 더 높아질 수 있는 좋은 기회였다...