[백준] 12899번 데이터 구조 C++ 코드

코드

#include <iostream>
#include <vector>
#include <cstring>
#include <map>
#include <cmath>
#include <algorithm>

using namespace std;
typedef long long ll;
const int MAX_N = 2000003;

ll n, a[MAX_N];
vector<ll> tree;

void update(vector<ll> &tree, int idx, ll diff) {
    while(idx <= MAX_N) {
        tree[idx] += diff;
        idx += (idx & -idx);
    }
}

ll sum(vector<ll> &tree, int idx) {
    ll ans = 0;
    while(idx > 0) {
        ans += tree[idx];
        idx -= (idx & -idx);
    }
    return ans;
}

int find_kth(int x) {
    int s = 1;
    int e = MAX_N;
    while(s < e) {
        int mid = (s + e) / 2;
        if(sum(tree, mid) >= x) e = mid;
        else s = mid + 1;
    }
    return s;
}

void FASTIO() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);
}

int main() {
    FASTIO();
    cin >> n;
    tree.resize(MAX_N + 1, 0);
    int num = 1;
    for(int i = 0; i < n; i++) {
        int b, c;
        cin >> b >> c;
        if (b == 1) {
            update(tree, c, 1);
        } else if (b == 2) {
            int kth_val = find_kth(c);
            cout << kth_val << '\n';
            update(tree, kth_val, -1);
        }
    }
    return 0;
}

풀이

사탕상자랑 비슷한 문제. x번째로 작은 수를 찾아 응답하고 그 수를 삭제하는 문제이다. 개수를 세어주면 되므로 펜윅 트리를 사용해서 구현했다. T가 1일 때는 트리에 넣고 개수를 증가시킨다. T가 2일 때는 n번째로 작은 수를 찾아서 삭제한다. 이 n번째 수를 찾을 때는 이분탐색을 이용한다. 인덱스에 개수가 저장되어있기 때문에 정렬된 상태라서 이분탐색을 사용할 수 있다. start와 end를 조정하며 n번째로 작은 수를 찾아준 후 그 숫자를 출력한다. 그리고 해당 인덱스에 -1을 해 삭제처리해주면 된다.