📌 문제
수 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으로 나누어 떨어지는 구간의 개수를 출력한다.
📌 문제 풀이
👨🏫 접근
구간합 문제의 변형 문제이다.
대체적으로 풀이 저변으로 구간합 알고리즘으로 구간 합 테이블을 응용해야 한다. 구간합을 빠르고 쉽게 구하기 위해서는 이전 값에 본인의 값을 더하여 테이블을 지속적으로 갱신해 만들어 주는 것이 좋다.
1, 2, 3, 4, 5, 6, 7, 8, 9, 10에서 6부터 9까지 더한다고 하는 것은, 1~9를 모두 더한 것에서 1~5까지 더한 것을 빼는 것과 같다. 여기서 1부터 10까지 이전의 값과 자신의 값을 더해보면 아래와 같다.
1, 3, 6, 10, 15, 21, 28, 36, 45, 55
여기서 9번(인덱스는 8) 숫자인 45와 5번(인덱스는 4번) 숫자인 15를 빼주면 30이며, 6부터 9까지 손수 더한 값도 30이다. 이처럼 미리 테이블에 이전 값들을 모두 더한 값을 저장해 두어서 원하는 구간 합을 빠르게 얻을 수 있다.
이를 나머지로 응용하면 된다.
1. 구간 합 테이블 만들기
여기서는 이전 값에 자신의 값을 더할 때 나누기 연산도 함께 진행했다. 즉, 연산을 하나 더 추가한 것으로 아래의 식과 같다.
prefix_sum[i] = (prefix_sum[i] - prefix_sum[i - 1]) % m
큰 수가 들어오더라도 사전에 값의 크기를 줄일 수 있었다(사실 나누지 않아도 복잡도 차이가 거의 미미했다).
2. 나머지 조건을 만족하는 구간 합 찾기
여기서는 구간 합의 나머지가 M 값과 같은 구간 합의 개수를 찾아야 한다. 그래서 가장 처음 시도한 것은 구간의 길이를 1부터 n까지 늘려가며 찾아보고 조건에 따라 카운팅을 추가하는 것이었다. 하지만, 이렇게 하면 시간 복잡도가 너무 컸다. 왜냐하면, 처음 구간 합 테이블을 구할 때부터 일단 O(n) 시간 복잡도가 소요된다.
그리고, 1부터 n까지의 반복하면서 리스트를 모두 순회하기에 거기서도 시간 복잡도가 많이 소요되었다. 일단 길이가 1이라는 것은 모든 원소를 돌아보는 것으로 n번의 작업이, 길이가 2면 n-1번의 작업, ... , 길이가 n이면 1번의 작업이 소요되므로 총 O(n(n + 1) / 2) ≈ O(n2)의 시간 복잡도가 소요된다. 그래서 둘을 더하면 O(n) + O(n2) ≈ O(n2) 이다.
💡 나머지가 같은 구간합 찾기
그래서 왜 나머지라는 조건을 넣었을까 생각을 해보면서 풀이 방법을 고민해봤다. 생각해보면 나머지가 같은 값들은 서로 빼거나 더해도 나머지가 같다. 분배법칙을 (123 % 3 +/- 456 % 3) % 3 = (123 +/- 456) % 3 이라는 모듈러 연산의 분배법칙이 성립하기 때문이다.
이를 구간 합으로 확장해봤다. 나머지 같은 구간합끼리 빼주고 나면 나머지가 0이 되기 때문에 이는 M으로 나눴을 때 나누어 떨어지는 구간 합이 된다. 예를 들어서,
10 4
4 2 8 7 1 5 9 2 4 3
위와 같은 입력이 들어왔다고 생각했을 때,
4 + 2 + 8 + 7 + 1 + 5 + 9 + 2를 4로 나눈 나머지는 38 % 4 = 2이다. 그리고 4 + 2 + 8 + 7 + 1를 4로 나눈 나머지는 22 % 4 = 2이다. 이때, 두 구간합을 빼주면 38 - 22 = 16이고 나머지는 0이 된다. 이를 풀어서 보자면, (4 + 2 + 8 + 7 + 1 + 5 + 9 + 2) - (4 + 2 + 8 + 7 + 1) = 5 + 9 + 2 = 16이라는 구간의 합이며 위 수열의 여러 구간들 중 하나이다.
따라서, 나머지가 같은 값들끼리 빼주고 나머지 연산을 하면 그 나머지는 0으로 나누어 떨어지게 되는 것이다.
3. 나머지 맵 만들기
그래서 나머지 맵을 만들어서 개수를 더했다. 어떤 수열에서 나머지가 k인 구간 합이 n개 있다는 것을 저장해두기 위해서이다. 나머지가 같은 값들은 1부터 n까지 모두 더해주면 된다. 이것은 직접 그 이유를 찾아보는 것이 도움이 될 것 같다. 간단히 생각해서 구간 합을 구하는 과정에서 나머지가 k인 맵에 n이라는 값이 매핑이 되어 있으면, 나중에 나머지가 k인 구간합을 발견했을 때, 이 구간 합은 n개의 구간합과 모두 연산을 해도 모두 나머지가 0이 될 수 있다는 의미로 최종적으로는 1부터 n까지 모두 값을 더해주는 것과 같다.
이렇게 하여 나머지 구간 합을 구할 수 있다. 이때, 나머지가 0인 구간합은 이미 그 자체로 조건을 충족시키기 때문에 0은 미리 1로 값을 초기화 해주고 시작해야 한다.
👨🏫 문제 풀이
📄 전체 코드
n, m = map(int, input().split())
arr = list(map(int, input().split()))
rest_count = { 0: 1 }
ans = 0
prefix_sum = [0] + arr
for i in range(1, n + 1):
prefix_sum[i] = r = (prefix_sum[i] + prefix_sum[i - 1]) % m
ans += rest_count.get(r, 0)
rest_count[r] = rest_count.get(r, 0) + 1
print(ans)
📄 준비
n, m = map(int, input().split())
arr = list(map(int, input().split()))
rest_count = { 0: 1 }
ans = 0
매핑을 위하여 딕셔너리 자료형을 사용했다.
📄 풀이
prefix_sum = [0] + arr
for i in range(1, n + 1):
prefix_sum[i] = r = (prefix_sum[i] + prefix_sum[i - 1]) % m
ans += rest_count.get(r, 0)
rest_count[r] = rest_count.get(r, 0) + 1
print(ans)
딕셔너리의 get 메서드는 첫번째 인자 값을 키로하여 값을 가져오며 만약 해당 키가 존재하지 않는다면, 두 번째 인자값을 기본값으로 반환한다. 그래서 이를 통해 쉽게 빈도수를 체크할 수 있다.
📌 총평
어려웠다... 나머지 구간을 어떻게 구해야 할까 생각해내는 게 쉽지 않은 과정이었다.