[백준 12899번] [세그먼트 트리] 데이터 구조

Approach

처음에는 세그먼트 트리 + 이분탐색으로 접근하여서 TLE가 나왔다.

T가 2인 쿼리가 들어오면 결과적으로 k번째 수를 찾아야 한다.

 

잘 생각해보면 이는 이분탐색으로 구할 수 있음을 쉽게 알 수 있다.

예를 들어, a라는 숫자가 k번째 수라고 가정했을 경우에는 1 ~ a - 1까지에 값에 존재하는 수는 k 미만이고, a이상의 수에 대해서는 k 이상이므로 구획이 정확하게 분할되게 된다. 따라서 이를 이분탐색을 활용해서 풀기 위해서는 1 ~ i까지에 존재하는 수의 개수만 관리해주면 된다.

 

여기까지 오면, prefix_sum을 생각할 수도 있고 segment tree를 생각할 수도 있다.

하지만, T가 1일때 값을 바꿔줘야하는 부분도 있기 때문에 segment tree를 사용하는 것이 더 유리하기는 하다.

 

위의 방식으로 문제를 제출하게 되면 TLE가 나오게 된다. 시간 복잡도를 계산했을 때 간당간당하길래 안될 것 같기는 했다..

 

근데, 잘 생각해보면 세그먼트 트리가 결과적으로 이분탐색의 느낌을 가지고 있다는 것을 쉽게 알 수 있다.

만약 특정한 쿼리에서 바라보고 있는 영역이 [start, end]이고, k번째 수를 찾고 있다고 하고, 세그먼트 트리는 구간 안에 존재하고 있는 숫자의 개수를 저장하였다고 가정하자. mid = (start + end) / 2라고 잡았을 때, [start, mid]에 존재하는 수가 k이상일 경우에는 [start, mid] 구간만 봐도 무방하다. 반면 미만일 경우에는 [mid + 1, end]에 해당하는 영역에서 k - seg([start, mid]) (단, seg[start, mid]는 [start, mid]에 해당하는 세그먼트 트리의 값)번째 값을 구해주면 된다.

 

세그먼트 트리가 본질적으로 이분탐색 기법을 활용하고 있다는 것을 학습하기 좋은 문제라고 생각한다.

Code

#include <bits/stdc++.h>
#define fastio cin.tie(0)->sync_with_stdio(0)

using namespace std;
using ll = long long;
using pii = pair<int, int>;
using tiii = tuple<int, int, int>;

int move_x[4] = {-1, 1, 0, 0};
int move_y[4] = {0, 0, -1, 1};
int seg[8000000];
void update(int node, int start, int end, int index, int value){
    if(index < start || end < index) return;
    seg[node] += value;
    if(start == end) return;
    int mid = (start + end) >> 1;
    update(node * 2, start, mid, index, value);
    update(node * 2 + 1, mid + 1, end, index, value);
}

int query(int node, int start, int end, int value){
    seg[node]--;
    if(start == end) return start;
    int mid = (start + end) >> 1;
    if(value <= seg[node * 2]) return query(node * 2, start, mid, value);
    else return query(node * 2 + 1, mid + 1, end, value - seg[node * 2]);
}


int main() {
    fastio;
    memset(seg, 0, sizeof(seg));
    int n;
    cin >> n;
    for(int i = 0; i < n; i++){
        int a, b;
        cin >> a >> b;
        if(a == 1){
            update(1, 1, 2000000, b, 1);
        }
        else{
            int val = query(1, 1, 2000000, b);
            cout << val << "\n";
        }
    }
    return 0;
}