배열과 여러 개의 쿼리가 주어집니다. 쿼리에는 두 가지 종류가 있습니다. update[L, R]는 L부터 R까지의 모든 요소를 각각의 제곱근 값으로 변경하는 연산이고, query[L, R]는 L부터 R까지의 요소 합을 계산하는 연산입니다. 여기서는 1-based 인덱스 배열을 사용한다고 가정합니다.
Input: nums[ ] = { 0, 9, 4, 1, 5, 2, 3 }, Query[ ] = { {1, 1, 3}, {2, 1, 2}, {1, 2, 5}, { 1, 4, 5}}
Output: 14
10
7
첫 번째 쿼리의 첫 번째 요소가 1이므로, 1부터 3까지의 구간 합을 계산합니다. 즉, 9 + 4 + 1 = 14입니다.
두 번째 쿼리의 첫 번째 요소가 2이므로, 1부터 2까지의 요소를 각각 제곱근 값으로 업데이트합니다. 새 배열은 { 3, 2, 1, 5, 2, 3 }이 됩니다.
세 번째 쿼리의 첫 번째 요소가 1이므로, 2부터 5까지의 구간 합을 계산합니다. 즉, 2 + 1 + 5 + 2 = 10입니다.
네 번째 쿼리의 첫 번째 요소가 1이므로, 4부터 5까지의 구간 합을 계산합니다. 즉, 5 + 2 = 7입니다.
Input: nums[] = { 0, 3, 2, 4, 16, 2 }, Query[ ] = {{1, 1, 3}, {2, 2, 5}}
Output: 9
문제 해결 접근 방식
단순한 접근 방식
모든 쿼리를 처음부터 끝까지 순회하면서, 합 쿼리에 대해서는 해당 구간의 합을 계산해 반환하고, 업데이트 쿼리에 대해서는 배열 요소를 직접 변경하는 방식으로 문제를 해결할 수 있습니다. 하지만 이 방법의 시간 복잡도는 O(q × n)으로, 배열의 크기와 쿼리 수가 커질수록 매우 비효율적입니다. 따라서 더 효율적인 접근 방식이 필요합니다.
효율적인 접근 방식
연산 횟수나 반복 횟수를 줄이면 프로그램의 성능을 크게 향상시킬 수 있습니다. 이때 활용할 수 있는 자료구조가 바로 이진 인덱스 트리(Binary Indexed Tree, 펜윅 트리)입니다. BIT는 업데이트와 구간 합 조회를 모두 O(log n) 시간에 처리할 수 있는 강력한 도구입니다.
핵심 아이디어는 다음과 같습니다. 업데이트 쿼리를 처리할 때 값이 이미 1인 요소는 제곱근을 취해도 그대로 1이므로, 다시 업데이트할 필요가 없습니다. 따라서 1보다 큰 값을 가진 인덱스들만 set(집합)에 저장해 두고, 이분 탐색(lower_bound)으로 구간의 시작 위치를 빠르게 찾은 뒤 순차적으로 업데이트를 진행합니다. 업데이트 결과 값이 1이 되었다면 해당 인덱스를 set에서 제거하여, 이후 쿼리에서 불필요한 갱신이 일어나지 않도록 합니다.
참고로 어떤 큰 수라도 제곱근 연산을 몇 차례(대체로 6회 이내) 반복하면 금방 1이 되기 때문에, 전체 업데이트 연산의 총 횟수는 제한적입니다. 이것이 이 접근법이 효율적으로 동작하는 핵심 이유입니다.
합 쿼리의 경우, query(R) − query(L−1) 공식을 이용하면 구간 [L, R]의 합을 간단하게 구할 수 있습니다.
예제 코드
위 접근 방식을 구현한 C++ 코드
#include <bits/stdc++.h>
using namespace std;
// 입력 배열의 최대 크기
const int m = 200;
// 이진 인덱스 트리 생성
int binary_indexed[m + 1];
// 업데이트 쿼리용 함수
void update_q(int a, int x, int n){
while(a <= n) {
binary_indexed[a] += x;
a += a & -a;
}
}
// 구간 합을 계산하는 함수
int sum_q(int a){
int s = 0;
while(a > 0) {
s += binary_indexed[a];
a -= a & -a;
}
return s;
}
int main(){
int no_query = 4;
int nums[] = { 0, 9, 4, 1, 5, 2, 3 };
int n = sizeof(nums) / sizeof(nums[0]);
// 쿼리를 저장하는 2차원 배열
int q[no_query + 1][3];
q[0][0] = 1, q[0][1] = 1, q[0][2] = 3;
q[1][0] = 2, q[1][1] = 1, q[1][2] = 2;
q[2][0] = 1, q[2][1] = 2, q[2][2] = 5;
q[3][0] = 1, q[3][1] = 4, q[3][2] = 5;
set<int> s;
for (int i = 1; i < n; i++) {
// 값이 1보다 큰 요소들의 인덱스를 set에 삽입
if (nums[i] > 1)
s.insert(i);
update_q(i, nums[i], n);
}
for (int i = 0; i < no_query; i++) {
// 첫 번째 값으로 업데이트 쿼리인지 합 쿼리인지 판별
if (q[i][0] == 2) {
while (true) {
// 이분 탐색으로 왼쪽 인덱스를 찾음
auto it = s.lower_bound(q[i][1]);
// 오른쪽 경계에 도달했는지 확인
if (it == s.end() || *it > q[i][2])
break;
q[i][1] = *it;
// 배열 요소를 제곱근 값으로 업데이트
update_q(*it, (int)sqrt(nums[*it]) - nums[*it], n);
nums[*it] = (int)sqrt(nums[*it]);
// 업데이트된 값이 1이라면 set에서 제거
if (nums[*it] == 1)
s.erase(*it);
q[i][1]++;
}
} else {
cout <<"query" << i+1 <<": " << (sum_q(q[i][2]) - sum_q(q[i][1] - 1)) << endl;
}
}
return 0;
}
실행 결과
query1: 14 query3: 10 query4: 7
마무리
이 튜토리얼에서는 배열에 대한 구간 합 쿼리와 제곱근 구간 업데이트 문제를 다루었습니다. 단순 반복문을 사용하는 기본 접근법의 한계를 살펴본 뒤, 이진 인덱스 트리(BIT)와 set을 결합해 효율성을 높이는 최적화 기법을 학습했습니다. 소개한 알고리즘은 C++뿐 아니라 C, Java, Python 등 다른 프로그래밍 언어로도 손쉽게 구현할 수 있으니, 직접 코드를 작성하며 개념을 익혀 보시기 바랍니다. 이 글이 여러분의 학습에 도움이 되기를 바랍니다.