洛谷P3369题解：Treap数据结构从入门到精通

9小时前洛谷题解34

一、数据结构概述

Treap是一种同时具备二叉搜索树(BST)和堆(Heap)性质的数据结构，通过随机优先级维护平衡性，实现高效的插入、删除和查询操作。

二、核心实现解析

节点结构：包含值(val)、计数(cnt)、子树大小(size)和随机优先级(priority)
旋转操作：左旋(rotateLeft)和右旋(rotateRight)维护堆性质
更新机制：update函数维护子树大小信息
六大核心操作：插入、删除、查排名、查值、前驱、后继

三、完整代码实现(带详细注释)

#include <iostream>
#include <cstdlib>
#include <climits>
using namespACe std;

const int INF = 1e8; // 处理大数值范围

struct Node {
    int val, size, cnt;
    int priority;
    Node *l, *r;
    Node(int v) : val(v), size(1), cnt(1), l(nullptr), r(nullptr) {
        priority = rand() % INF; // 随机优先级保证平衡
    }
};

class Treap {
private:
    Node *root;
    
    // 更新节点子树大小
    void update(Node *node) {
        if(!node) return;
        node->size = node->cnt;
        if(node->l) node->size += node->l->size;
        if(node->r) node->size += node->r->size;
    }
    
    // 左旋操作
    void rotateLeft(Node *&node) {
        Node *temp = node->r;
        node->r = temp->l;
        temp->l = node;
        update(node);
        update(temp);
        node = temp;
    }
    
    // 右旋操作
    void rotateRight(Node *&node) {
        Node *temp = node->l;
        node->l = temp->r;
        temp->r = node;
        update(node);
        update(temp);
        node = temp;
    }
    
    // 插入操作
    void insert(Node *&node, int val) {
        if(!node) {
            node = new Node(val);
            return;
        }
        if(val == node->val) {
            node->cnt++; // 重复值计数
        } else if(val < node->val) {
            insert(node->l, val);
            if(node->l->priority > node->priority)
                rotateRight(node); // 维护堆性质
        } else {
            insert(node->r, val);
            if(node->r->priority > node->priority)
                rotateLeft(node); // 维护堆性质
        }
        update(node);
    }
    
    // 删除操作
    void remove(Node *&node, int val) {
        if(!node) return;
        if(val < node->val) {
            remove(node->l, val);
        } else if(val > node->val) {
            remove(node->r, val);
        } else {
            if(node->cnt > 1) {
                node->cnt--; // 减少计数
            } else {
                if(!node->l || !node->r) {
                    Node *temp = node->l ? node->l : node->r;
                    delete node;
                    node = temp; // 单子树情况直接替换
                } else {
                    // 选择优先级高的子树旋转
                    if(node->l->priority > node->r->priority) {
                        rotateRight(node);
                        remove(node->r, val);
                    } else {
                        rotateLeft(node);
                        remove(node->l, val);
                    }
                }
            }
        }
        if(node) update(node);
    }
    
    // 获取排名
    int getRank(Node *node, int val) {
        if(!node) return 0;
        if(val < node->val) return getRank(node->l, val);
        int leftSize = node->l ? node->l->size : 0;
        if(val == node->val) return leftSize + 1;
        return leftSize + node->cnt + getRank(node->r, val);
    }
    
    // 根据排名获取值
    int getValue(Node *node, int rank) {
        if(!node) return INF;
        int leftSize = node->l ? node->l->size : 0;
        if(rank <= leftSize) return getValue(node->l, rank);
        if(rank <= leftSize + node->cnt) return node->val;
        return getValue(node->r, rank - leftSize - node->cnt);
    }
    
    // 获取前驱
    int getPre(Node *node, int val) {
        if(!node) return -INF;
        if(node->val >= val) return getPre(node->l, val);
        return max(node->val, getPre(node->r, val));
    }
    
    // 获取后继
    int getNext(Node *node, int val) {
        if(!node) return INF;
        if(node->val <= val) return getNext(node->r, val);
        return min(node->val, getNext(node->l, val));
    }

public:
    Treap() : root(nullptr) { srand(time(0)); }
    
    // 公开接口
    void insert(int val) { insert(root, val); }
    void remove(int val) { remove(root, val); }
    int getRank(int val) { return getRank(root, val); }
    int getValue(int rank) { return getValue(root, rank); }
    int getPre(int val) { return getPre(root, val); }
    int getNext(int val) { return getNext(root, val); }
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(0);
    
    Treap treap;
    int n, opt, x;
    cin >> n;
    while(n--) {
        cin >> opt >> x;
        switch(opt) {
            case 1: treap.insert(x); break;
            case 2: treap.remove(x); break;
            case 3: cout << treap.getRank(x) << '\n'; break;
            case 4: cout << treap.getValue(x) << '\n'; break;
            case 5: cout << treap.getPre(x) << '\n'; break;
            case 6: cout << treap.getNext(x) << '\n'; break;
        }
    }
    return 0;
}