백준 12899 데이터 구조 혼내주기

 

 

 

세그 트리로 k번째 수 찾기.
입력되는 숫자의 범위가 정해져 있다는 사실을 이용해 세그 트리의 leaf node가 각 수의 갯수를 나타내고 parent node들은 자식 노드들의 갯수의 합을 나타내도록 구현하면 된다. query는 binary search처럼 k번째 수를 찾으면 된다.
k번째도 그렇고 머지 소트 트리도 그렇고 주어진 입력에서도 많은 정보를 얻을 수 있다. 그저 내 꼼꼼함이 부족할 뿐...

 

 

#include <bits/stdc++.h>
#define fastio ios::sync_with_stdio(0), cin.tie(0)
using namespace std;

int tree[1<<22];
int n;

void input() {
    fastio;
    cin >> n;
}

int update(int t, int v, int i = 1, int s = 0, int e = (1<<21)-1) {
    if(t < s || t > e) return tree[i];
    if(s == e) return tree[i] += v;
    int m = (s + e) >> 1;
    return tree[i] = update(t, v, i<<1, s, m) + update(t, v, i<<1|1, m+1, e);
}

int query(int k, int i = 1, int s = 0, int e = (1<<21)-1) {
    if(s == e) return s;
    int m = (s + e) >> 1;
    if(k <= tree[i<<1]) return query(k, i<<1, s, m);
    return query(k-tree[i<<1], i<<1|1, m+1, e);
}

int main() {
    input();
    while(n--) {
        int a, b; cin >> a >> b;
        if(a == 1) {
            update(b-1, 1);
        }
        else {
            int idx = query(b);
            cout << idx+1 << "\n";
            update(idx, -1);
        }
    }
}