문제 개요
크기가 n인 배열 nums에 양의 정수들이 들어 있다고 가정해 보겠습니다. 그리고 정수 쌍 (pi, qi)를 담고 있는 또 다른 배열 queries가 주어집니다. queries 배열의 각 쿼리에 대한 답은, pi <= j < n이면서 (j - pi)가 qi로 나누어 떨어지는 모든 nums[j] 값들의 합입니다. 모든 쿼리에 대한 답을 반환해야 하며, 계산 결과가 너무 커질 경우에는 10^9 + 7로 나눈 나머지를 반환합니다.
예를 들어, 입력이 다음과 같다고 해보겠습니다.
nums = [2, 3, 4, 5, 6, 7, 8, 9, 10] queries = [(2, 5), (7, 3), (6, 4)]
이 경우 출력 결과는 [13, 9, 8]이 됩니다.
해결 접근 방식
이 문제는 제곱근 분할(sqrt decomposition) 아이디어를 활용해 효율적으로 해결할 수 있습니다. 간격(k)이 작은 쿼리는 미리 계산해 둔 누적합 테이블에서 즉시 답을 가져오고, 간격이 큰 쿼리는 해당 구간의 원소 수가 적으므로 직접 더해서 처리합니다.
구체적인 단계는 다음과 같습니다.
- A := nums
- Q := queries
- n := nums의 길이
- M := 10^9 + 7
- m := int(n ** 0.5) + 2
- P := 리스트 A를 m번 복사해 만든 새로운 2차원 리스트
- i를 1부터 m-1까지 반복하며:
- j를 n-1부터 0까지 역순으로 반복하며:
- 만약 i + j < n이라면, P[i][j] := (P[i][j] + P[i][i+j]) % M 로 갱신합니다. 이 과정을 거치면 P[i][j]는 인덱스 j부터 시작해 간격 i마다 뒤따르는 모든 원소의 합이 됩니다.
- j를 n-1부터 0까지 역순으로 반복하며:
- Q의 각 쌍 (b, k)에 대해:
- 만약 k < m이면, 미리 계산된 P[k][b] 값을 반환합니다.
- 그렇지 않으면, sum(A[b::k]) % M 을 직접 계산해 반환합니다.
예제 코드
아래 파이썬 구현을 통해 동작 방식을 더 명확히 이해할 수 있습니다.
def solve(A, Q):
n, M = len(A), 10**9+7
m = int(n**0.5)+2
P = [A[:] for _ in range(m)]
for i in range(1,m):
for j in range(n-1,-1,-1):
if i+j < n:
P[i][j] = (P[i][j]+P[i][i+j]) % M
return [P[k][b] if k < m else sum(A[b::k]) % M for b, k in Q]
print(solve([2, 3, 4, 5, 6, 7, 8, 9, 10], [(2, 5), (7, 3), (6, 4)]))입력
[2, 3, 4, 5, 6, 7, 8, 9, 10], [(2, 5), (7, 3), (6, 4)]
출력
[13, 9, 8]
정리
이 알고리즘은 전처리 단계에서 O(m × n), 즉 약 O(n√n) 시간이 소요되지만, 이후 간격이 작은 쿼리는 O(1)에 즉시 답변할 수 있습니다. 간격이 큰 쿼리도 한 번 확인할 때 최대 √n개 정도의 원소만 더하면 되므로 전체적으로 효율적인 성능을 보장합니다. 대량의 범위 합 쿼리를 빠르게 처리해야 하는 상황에서 유용하게 활용할 수 있는 패턴입니다.