每日一题-16th

⭐每日一题⭐专栏

written by SJTU-XHW

本人学识有限，解析难免有错，恳请读者能够批评指正，本人将不胜感激！

LuoGu P5076. “二叉树”

题目描述

您需要写一种数据结构，来维护一些数（都是 $10^9$ 以内的数字）的集合，最开始时集合是空的。其中需要提供以下操作，操作次数 $q$ 不超过 $10^4$：

查询 $x$ 数的排名（排名定义为比当前数小的数的个数 $+1$。若有多个相同的数，应输出最小的排名）。
查询排名为 $x$ 的数。
求 $x$ 的前驱（前驱定义为小于 $x$，且最大的数）。若未找到则输出 $-2147483647$。
求 $x$ 的后继（后继定义为大于 $x$，且最小的数）。若未找到则输出 $2147483647$。
插入一个数 $x$。

说明

输入格式

第一行是一个整数 $q$，表示操作次数。

接下来 $q$ 行，每行两个整数 $op,x$，分别表示操作序号以及操作的参数 $x$。

输出格式

输出有若干行。对于操作 $1,2,3,4$，输出一个整数，表示该操作的结果。

样例 1 输入

样例 1 输出

思路和实现

首先看要求，“查询 $x$ 数的排名”，我寻思着普通的二叉树好像不太行，时间复杂度不如存成线性结构；而且平衡二叉树（AVL 和红黑树）好像也不太行，毕竟是求排名；

唯一比较靠谱的是二叉排序树（即二叉查找树），因为它的中序序列就是排序后的序列；于是决定先选择二叉搜索树（元素全为 int，$10^9$ 以内）：

#include <bits/stdc++.h>
using namespace std;

// 在已排序数组中使用二分查找
int binSearch(int* ord, int left, int right, int element, bool showLc = 0) {
    int mid;
    while (left < right) {
        mid = (left + right) / 2;
        if (ord[mid] < element) left = mid + 1;
        else if (ord[mid] > element) right = mid - 1;
        else return mid;
    }
    if (left == right && ord[left] == element) return left;
    return showLc ? left : -1;
}

class BSTree {
private:
    struct node {
        int data;
        node* left;
        node* right;
        node(int dt = 0, node* L = 0, node* R = 0)
            : data(dt), left(L), right(R) {}
    };
    
    int size;
    node* root;
    // 懒加载
    bool updated;
    int* orderedArr;

    // 中序遍历的递归实现
    void midOrderTrav();
    void midOrderTrav(node* curRoot, int& idx);
    void insert(node** curRoot, int x) {
        while (*curRoot) {
            if ((*curRoot)->data >= x) curRoot = &((*curRoot)->left);
            else curRoot = &((*curRoot)->right);
        }
        *curRoot = new node(x);
    }
    void clear(node*& curRoot) {
        if (curRoot == nullptr) return;
        clear(curRoot->left);
        clear(curRoot->right);
        delete curRoot; curRoot = nullptr;
    }
    
public:
    BSTree(): size(0), root(0), updated(1), orderedArr(0) {}
    ~BSTree() { clear(root); }
    void updateArr() {
        if (size == 0) return;
        midOrderTrav();
        updated = 1;
    }
    void insert(int x) {
        updated = 0; ++size;
        insert(&root, x);
    }
    int getPre(int x) {
        if (!updated) updateArr();
        int idx = binSearch(orderedArr, 0, size - 1, x);
        if (idx == -1 || idx == 0) return -2147483647;
        return orderedArr[idx - 1];
    }
    int getPost(int x) {
        if (!updated) updateArr();
        int idx = binSearch(orderedArr, 0, size - 1, x);
        if (idx == -1 || idx == size - 1) return 2147483647;
        return orderedArr[idx + 1];
    }
    int getRank(int x) {
        if (!updated) updateArr();
        int idx = binSearch(orderedArr, 0, size - 1, x, true);
        if (orderedArr[idx] < x) return idx + 2;
        else return idx + 1;
    }
    int getByRank(int rank) {
        if (!updated) updateArr();
        return orderedArr[rank - 1];
    }
};

void BSTree::midOrderTrav() {
    if (orderedArr) delete[] orderedArr;
    if (size == 0) { orderedArr = 0; return; }
    orderedArr = new int[size] {0};
    int idx = 0;
    midOrderTrav(root, idx);
}

void BSTree::midOrderTrav(BSTree::node* curRoot, int& idx) {
    if (idx >= size || curRoot == nullptr) return;
    midOrderTrav(curRoot->left, idx);
    if (idx < size) {
        orderedArr[idx++] = curRoot->data;
        midOrderTrav(curRoot->right, idx);
    }
}


int main() {
    BSTree tree; int cnt, op, x; scanf("%d", &cnt);
    for (int i = 0; i < cnt; ++i) {
        scanf("%d%d", &op, &x);
        switch (op) {
            case 5: tree.insert(x); break;
            case 4: printf("%d\n", tree.getPost(x)); break;
            case 3: printf("%d\n", tree.getPre(x)); break;
            case 2: printf("%d\n", tree.getByRank(x)); break;
            default: printf("%d\n", tree.getRank(x));
        }
    }
    return 0;
}