Treap树的C++类实现

Treap:

一种数据结构，结合了二叉搜索树和堆的特点，为了避免BST树插入有序数列时一边树过高的问题，引入了随机权值，在插入时进行随机权值的比较来进行左旋右旋的操作，进而达到平衡的效果，目的与avl类似，但操作方便，代码更加简洁。支持插入节点、删除节点、求第x大的节点、求权值为x的节点的排名、求权值比x小的最大节点、求权值比x大的最小节点

除建树外，各操作时间复杂度均为O(logn)

#include<cstdio>
#include<cstdlib>
#include<climits>
#include<cmath>
#include<iostream>

using namespace std;

struct Node {
    int data;
    Node *left;
    Node *right;
    int priority;
    int size;
    int cnt;
    Node(int value, int level) : data(value), left(NULL), right(NULL), priority(level), size(1), cnt(1) {}
};

class Treap {
private:
    Node *root;
    void leftRotate(Node* &p);                  //左旋
    void rightRotate(Node* &p);                 //右旋
    void insert(Node* &p, int value);           //插入节点的递归写法
    void remove(Node* &p, int value);           //删除节点的递归写法
    int get_rank_by_value(Node* &p, int value); //通过排名获取节点值的递归写法
    int get_value_by_rank(Node* &p, int rank);  //通过节点值获取排名的递归写法
    int get_prev(Node* &p, int value);          //找到节点的前驱的递归写法
    int get_next(Node* &p, int value);          //找到节点的后继的递归写法
    void pushup(Node* &p);                      //更新节点的参数
public:
    Treap();
    void insert(int value);
    void remove(int x);
    bool search(int x);
    int get_rank_by_value(int value);
    int get_value_by_rank(int rank);
    int get_prev(int value);
    int get_next(int value);
    bool empty();                               //判断是否树为空
};

Treap::Treap() {
    root = NULL;
}

void Treap::pushup(Node* &p){
    if(!p) return;
    p->size = 0;
    if(p->left) p->size += p->left->size;
    if(p->right) p->size += p->right->size;
    p->size += p->cnt;
}

void Treap::leftRotate(Node* &p) {
    Node *k = p->right;
    p->right = k->left;
    k->left = p;
    p = k;
    pushup(p->left);
    pushup(p);
}

void Treap::rightRotate(Node* &p) {
    Node *k = p->left;
    p->left = k->right;
    k->right = p;
    p = k;
    pushup(p->right);
    pushup(p);
}

void Treap::insert(int value) {
    insert(root, value);
}

void Treap::insert(Node* &p, int value) {
    if (p == NULL) {
        p = new Node(value, rand());
    }
    else if (value == p->data) p->cnt++;
    else if (value < p->data) {
        insert(p->left, value);
        if(p->left && p->left->priority > p->priority) rightRotate(p);
    }
    else{
        insert(p->right, value);
        if(p->right && p->right->priority < p->priority) leftRotate(p);
    }
    
    pushup(p);
}

void Treap::remove(int value) {
    remove(root, value);
}

void Treap::remove(Node* &p, int value) {
    if(!p){cout<<"invalid value!"<<endl; return;}
    if (p->data == value) {
        if (p->left == NULL) {
            p = p->right;
        } else if (p->right == NULL) {
            p = p->left;
        } else {
            if (p->left->priority > p->right->priority) {
                rightRotate(p);
                remove(p->right, value);
            } else if (p->left->priority < p->right->priority) {
                leftRotate(p);
                remove(p->left, value);
            }
        }
    } else {
        if (value < p->data) {
            remove(p->left, value);
        } else {
            remove(p->right, value);
        }
    }
    pushup(p);
}

bool Treap::search(int value) {
    Node *p = root;
    while (p) {
        if (p->data == value) {
            return true;
        } else {
            p = p->data < value ? p->right : p->left;
        }
    }
    return false;
}

bool Treap::empty(){
    return root == nullptr ? true : false;
}

int Treap::get_rank_by_value(Node* &p, int value){
    if(!p) return -1;
    int l = 0;
    if(p->left) l = p->left->size;
    if(p->data == value) return l + 1;
    if(p->data > value) return get_rank_by_value(p->left, value);
    return l + p->cnt + get_rank_by_value(p->right, value);
}

int Treap::get_rank_by_value(int value){
    return get_rank_by_value(root, value);
}

int Treap::get_value_by_rank(Node* &p, int rank){
    if(!p) return INT_MAX;
    int l = 0;
    if(p->left) l = p->left->size;
    if(l >= rank) return get_value_by_rank(p->left, rank);
    if(l + p->cnt >= rank) return p->data;
    return get_value_by_rank(p->right, rank - l - p->cnt);
}

int Treap::get_value_by_rank(int rank){
    return get_value_by_rank(root, rank);
}

int Treap::get_prev(Node* &p, int value){
    if(!p) return -INT_MAX;
    if(p->data >= value) return get_prev(p->left, value);
    return max(p->data, get_prev(p->right, value));
}

int Treap::get_prev(int value){
    return get_prev(root, value);
}

int Treap::get_next(Node* &p, int value){
    if(!p) return INT_MAX;
    if(p->data <= value) return get_next(p->right, value);
    return min(p->data, get_next(p->left, value));
}

int Treap::get_next(int value){
    return get_next(root, value);
}

测试代码：

int main(int argc, char const *argv[]) {
    Treap treap;
    int N;
    cout<<"input len: ";
    cin>>N;
    cout<<endl;
    cout<<"input values: "<<endl;
    for (int i = 0; i < N; i++) {
        int value;
        cin >> value;
        treap.insert(value);
    }
    cout<<endl;
    cout<<endl;
    for (int i = 0; i < N; i++) {
        if(treap.empty()) break;
        int query;
        cout<<"input query: "<<endl;
        cin>>query;
        cout<<"value "<<query<<" rank is : "<<treap.get_rank_by_value(query)<<endl;
        cout<<"rank "<<query<<" value is : "<<treap.get_value_by_rank(query)<<endl;
        cout<<query<<" prev is : "<<treap.get_prev(query)<<endl;
        cout<<query<<" next is : "<<treap.get_next(query)<<endl;
        //treap.remove(query);
        cout<<endl;
    }
    return 0;
}

测试结果：