python/TangDou/AcWing/BalanceTree/TODO/253_Trie.cpp

#include <bits/stdc++.h>
using namespace std;
//	571 ms
const int N = 3200100;
const int base = 1e7;

// 01Trie树
int tr[N][2];
int root = 1; //根是1号点
int idx = 1;  //节点号计数器
int cnt[N];   //结点p的子结点的个数

//插入
void insert(int x, int t) {
    x += base; //保证输入非负
    int p = root;
    for (int i = 31; ~i; i--) {          //从高位到低位
        int c = x >> i & 1;              //计算x在第i位是0还是1
        if (!tr[p][c]) tr[p][c] = ++idx; //如果想走的位置不存在，则创建之,++idx为新创建的节点的编号
        p = tr[p][c];                    //游标p继续前行
        //这句是迎合本题特殊加入的记录数据
        cnt[p] += t; //以p为根的子树中节点个数
    }
}

//获取小于数x的个数
int rnk(int x) {
    x += base; //保证输入非负
    int p = root;
    int res = 0;
    for (int i = 31; ~i; i--) {      //从高位到低位
        int c = x >> i & 1;          //计算x在第i位是0还是1
        if (c) res += cnt[tr[p][0]]; //当该位为1时，找其他的该位为0的数的个数，就是严格小于这个数的个数了
        p = tr[p][c];                //游标p继续前行
    }
    return res;
}

//获取第k位置的数字
int find_kth(int k) {
    int p = root;
    int res = 0;
    for (int i = 31; ~i; i--) { // ~i 表示  i>=0
        if (k > cnt[tr[p][0]])  //在右手边
            res += (1 << i), k -= cnt[tr[p][0]], p = tr[p][1];
        //(1) res += (1 << i) //恢复原始数据的拼接过程
        //(2) x -= cnt[tr[p][0]] 减去左手边的个数
        //(3) p = tr[p][1] 向右手边继续游标前进
        else
            p = tr[p][0]; //游标p向左移动进去，继续前行
    }
    return res - base; //因为我们之前加过一个base，所以这里要减去
}

int main() {
    int n;
    scanf("%d", &n);
    while (n--) {
        int op, x;
        scanf("%d%d", &op, &x);
        if (op == 1) insert(x, 1);                             //插入x
        if (op == 2) insert(x, -1);                            //删除x
        if (op == 3) cout << rnk(x) + 1 << endl;               //查询严格小于x的个数，再+1，就是x的排名
        if (op == 4) cout << find_kth(x) << endl;              //查询第k位置的数
        if (op == 5) cout << find_kth(rnk(x)) << endl;         // x的前驱(小于x的最大值)
        if (op == 6) cout << find_kth(rnk(x + 1) + 1) << endl; // x的后继(大于x的最小值)
    }
    return 0;
}