백준[12899] 데이터 구조

문제 설명

문제 링크

http://icpc.me/12899

12899번: 데이터 구조

풀이

세그먼트 트리를 이용하여 k번째 수를 가져오는 문제이다.

세그먼트 트리를 만들 때 수의 범위가 아닌 수의 값으로 범위를 정해야 한다.

예를 들어, 5번째 수를 구할 때, 1 ~ N/2에 수가 6개 N/2 + 1 ~ N에 수가 3개가 있다고 한다면 왼쪽 노드로 이동하여 같은 작업을 반복해야 한다. 반면에 1 ~ N/2에 수가 2개 N/2 + 1 ~ N에 수가 5개가 있다고 하면, 오른쪽 노드로 이동하여 5-2인 3번째 수를 구해야 한다. 

코드

#include <cstdio>

const int MAX = 2e6;

int tree[MAX * 4 + 6];

int query(int node, int s, int e, int k) {
    tree[node]--;
    if (s == e) return s;
    int mid = (s + e) >> 1; 
    if (tree[2 * node] >= k) return query(2 * node, s, mid, k);
    else return query(2 * node + 1, mid + 1, e, k - tree[2 * node]);
}

void update(int node, int s, int e, int idx) {
    if (idx < s || e < idx) return;
    tree[node]++;
    if (s == e) return;
    int mid = (s + e) >> 1;
    update(2 * node, s, mid, idx);
    update(2 * node + 1, mid + 1, e, idx);
}

int main() {
    int n; scanf("%d", &n);
    while (n--) {
        int s, x; scanf("%d %d", &s, &x);
        if (s == 1) {
            update(1, 1, MAX, x);
        } else {
            printf("%d\n", query(1, 1, MAX, x));
        }
    }
}