python/TangDou/LuoGuBook/P5076.cpp

#include <bits/stdc++.h>

using namespace std;

//题解
//https://www.cnblogs.com/do-while-true/p/13566274.html
//const int INF = 2147483647;  这个 2147483647的十六进制表示：0x7fffffff,是等价的
const int INF = 0x7fffffff;
const int N = 1000010;
int cnt;
struct Node {
    int left;   //左儿子
    int right;  //右儿子
    int size;   //以该节点为根结点的子树的结点个数
    int value;  //该结点的权值
    int sq_sum;    //该结点权值出现的次数
} tree[N];
int q, opt, x;
/**
* 测试用例：
7
5 1
5 3
5 5
1 3
2 2
3 3
4 3

参考答案：
2
3
1
5
 */
//查询数x的排名(找出以root为根的树中，有多少个结点值小于查询数x)
int queryVal(int root, int x) {
    //root==0表示进入到一个结点，准备再向下，结果再向下没有了，就是：表示找不到权值为x的数，返回0
    if (root == 0) return 0;

    //如果x等于树根的value值，那么比它小的有tree[tree[root].left].size个，
    // 它是第tree[tree[root].left].size+1个
    if (x == tree[root].value) return tree[tree[root].left].size;

    //如果x小于树根的value值，那么需要到左子树中去计算排名
    if (x < tree[root].value) return queryVal(tree[root].left, x);

    //如果x大于树根的value值，那么需要到右子树中去计算排名,同时，需要加上左子树的
    // 结点个数+root的结点个数
    return queryVal(tree[root].right, x) + tree[tree[root].left].size + tree[root].num;
}

//查询排名为x的数
int queryRank(int root, int x) {
    //root==0表示进入到一个结点，准备再向下，结果再向下没有了，就是：表示找不到排名为x的数，返回0
    if (root == 0) return 0;//其实这里的有坑的，找不到题目没有说明是返回INF，还是返回0，
    // 实测返回INF也是可以AC的。

    //如果左子树中包含这个排名，那么递归到左子树去找
    if (tree[tree[root].left].size >= x) return queryRank(tree[root].left, x);
    //如果根结点的排名==x,那么返回根结点的权值
    if (tree[tree[root].left].size + tree[root].num >= x) return tree[root].value;
    //如果右子树中包含这个排名，那么递归到右子树中去找，这个子排名就是减去左子树的总结点个数
    // 再减去root的个数
    return queryRank(tree[root].right, x - tree[tree[root].left].size - tree[root].num);
}

//向二叉搜索树中增加一个权值为x的数字
void insert(int root, int x) {
    //由于增加了一个结点，root的总结点个数+1
    //这一步是递归的必要步骤，表示以它为根的家族中增加了一个结点
    tree[root].size++;

    //如果新增加的数字x与root根的权值相等，那么num++
    if (tree[root].value == x) {
        tree[root].num++;
    } else if (tree[root].value > x) {//如果root的权值大于x,左子树
        //如果左子树不空，则递归向左子树插入结点x
        if (tree[root].left != 0) insert(tree[root].left, x);
        else {
            //如果左子树为空，那么需要创建一个新结点
            tree[++cnt].value = x; //新的数组位置用来存储x
            tree[cnt].size = 1;    //新增加的结点，下面没有子树，所以size=1
            tree[cnt].num = 1;     //权值为x的目前只有1个
            tree[root].left = cnt; //添加为左子树
        }
    } else { //如果root的权值小于x，右子树
        //如果右子树不空，则递归向右子树插入结点x
        if (tree[root].right != 0) insert(tree[root].right, x);
        else {
            //如果右子树为空，那么需要创建一个新结点
            tree[++cnt].value = x;          //新的数组位置用来存储x
            tree[cnt].size = 1;             //新增加的结点，下面没有子树，所以size=1
            tree[cnt].num = 1;              //权值为x的目前只有1个
            tree[root].right = cnt;         //添加为右子树
        }
    }
}

/**
 * 功能：查询前驱
 * @param root 以root为根的树
 * @param x    要查找的数字x
 * @param ans  暂存的前驱值，待更新成最合理的前驱值
 * @return     数字x的前驱
 */
int queryPrev(int root, int x, int ans) {
    //如果根结点值大于x,那么需要去左子树中去找
    if (tree[root].value >= x) {
        //如果左子树为空的,将前面获取到的最优解返回就可以了
        if (tree[root].left == 0) return ans;
        else
            //到左子树中去查找
            return queryPrev(tree[root].left, x, ans);
    } else {
        //如果根结点值小于x,那么需要到右子树中去找
        //如果右子树为空，那么就是根结点的value值
        if (tree[root].right == 0) return tree[root].value;
        //右子树不为空，递归到右子树去查找
        return queryPrev(tree[root].right, x, tree[root].value);
    }
}

/**
 * 功能：查询后继
 * @param root 以root为根的树
 * @param x    要查找的数字x
 * @param ans  暂存的后继值，待更新成最合理的后继值
 * @return     数字x的后继
 */
int queryNext(int root, int x, int ans) {
    //如果根结点的权值小于等于x,需要向右子树去找
    if (tree[root].value <= x) {
        //如果右子树为空，则没有比它更大的，返回暂存的最合理值ans
        if (tree[root].right == 0) return ans;
        else
            //如果右子树不空，则递归到右子树中去找
            return queryNext(tree[root].right, x, ans);
    } else {
        //向左子树中去找

        //如果左子树为空，那么就是根的value值
        if (tree[root].left == 0) return tree[root].value;
        //如果左子树不空，递归到左子树去找
        return queryNext(tree[root].left, x, tree[root].value);
    }
}


int main() {
    //q次询问
    cin >> q;
    while (q--) {
        //操作命令和操作的值
        cin >> opt >> x;
        switch (opt) {
            case 1: //查询x数的排名,注意这后面的+1，其实,rnk就是在根为root的子树中，
            // 找到比x值小的结点的个数，然后再加1，就是x的排名
                cout << queryVal(1, x) + 1 << endl;
                break;
            case 2: //查询排名为x的数
                cout << queryRank(1, x) << endl;
                break;
            case 3: //求x的前驱
                cout << queryPrev(1, x, -INF) << endl;
                break;
            case 4: //求x的后继
                cout << queryNext(1, x, INF) << endl;
                break;
            case 5: //插入一个数 x
                if (cnt == 0) {                 //一个都没有，那么，第一个是根结点
                    cnt++;                      //数组中1号位置
                    tree[cnt].num = 1;          //权值是x的数量目前是1个
                    tree[cnt].size = 1;         //它+子树结点的总结点个数是1个
                    tree[cnt].value = x;        //权值是x
                } else insert(1, x);       //将数x插入到根为1的二叉搜索树中
                break;
        }
    }
    return 0;
}