P3369 【模板】普通平衡树

原创于 2025-12-21 12:22:22 发布 · 95 阅读

3 ·

CC 4.0 BY-SA版权

文章标签：

#AVL #平衡二叉树 #平衡二叉树插入 #平衡二叉树删除 #前驱 #后继 #指针

P3369 【模板】普通平衡树

题目描述

您需要动态地维护一个可重集合 $M$ ，并且提供以下操作：

向 $M$ 中插入一个数 $x$ 。
从 $M$ 中删除一个数 $x$ （若有多个相同的数，应只删除一个）。
查询 $M$ 中有多少个数比 $x$ 小，并且将得到的答案加一。
查询如果将 $M$ 从小到大排列后，排名位于第 $x$ 位的数。
查询 $M$ 中 $x$ 的前驱（前驱定义为小于 $x$ ，且最大的数）。
查询 $M$ 中 $x$ 的后继（后继定义为大于 $x$ ，且最小的数）。

对于操作 $3, 5, 6$ ，不保证当前可重集中存在数 $x$ 。

对于操作 $5, 6$ ，保证答案一定存在。

输入格式

第一行为 $n$ ，表示操作的个数，下面 $n$ 行每行有两个数 $opt\text{opt}$ 和 $x$ ， $opt\text{opt}$ 表示操作的序号（$ 1 \leq \text{opt} \leq 6 $）。

输出格式

对于操作 $3, 4, 5, 6$ 每行输出一个数，表示对应答案。

输入输出样例 #1

输入 #1

输出 #1

106465
84185
492737

说明/提示

【数据范围】

对于 $100%100\%$ 的数据， $1≤n≤1051\le n \le 10^5$ ， $\le 10^7$ 。

来源：Tyvj1728，原名：普通平衡树。

在此鸣谢！

代码

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

struct node{//节点结构体 
    int d,         // 节点存储的值
        height,    // 当前节点作为根的子树高度
        cnt,       // 该值的重复次数（相同值合并到一个节点）
        size,      // 当前子树的总元素数（=自身cnt+左子size+右子size）
        x;         // 平衡因子（左子树高度 - 右子树高度）
    node *f, *l, *r;// 父节点、左子节点、右子节点指针 
    // 构造函数：初始化新节点，默认高度/计数/元素数为1，平衡因子0，指针全空
    node(int val):d(val),height(1),cnt(1),size(1),x(0){f=l=r=nullptr;}
    void update(){// 更新节点的size、height、x（核心属性）
        size = cnt; // 先初始化size为自身重复数
        size += (l ? l->size : 0) + (r ? r->size : 0);// 加上左、右子树的总元素数
        int hl = (l ? l->height : 0), hr = (r ? r->height : 0);// 计算左右子树高度（空节点高度为0）
        height = max(hl, hr) + 1; // 子树高度=左右子树最高值+自身
        x = hl - hr; // 平衡因子=左高-右高（>1左失衡，<-1右失衡）
    }
};

class avltree{// AVL树类（自平衡二叉搜索树）
private:// 私有属性和方法（外部不可访问）
    node *root; // 树的根节点指针
    
    // 递归插入元素d：核心逻辑=递归找位置+回溯更新平衡
    node* insert(node *fa, int d){
        if(!fa) return new node(d); // 递归终止层：找到空位置，新建节点替代空指针（仅执行1次）
        if(d < fa->d){// 插入值更小，往左子树插
            fa->l = insert(fa->l, d); // 逐层递归，把新节点插入左子树
            if(fa->l) fa->l->f = fa;  // 绑定父子指针：仅倒数第二层有效（新节点的父节点层），其他层冗余但无害
        }else if(d > fa->d){// 插入值更大，往右子树插
            fa->r = insert(fa->r, d); // 逐层递归，把新节点插入右子树
            if(fa->r) fa->r->f = fa;  // 绑定父子指针：同上，仅倒数第二层有用
        }else{// 插入值已存在，仅更新计数（不新增节点）
            fa->cnt++;    // 重复数+1
            fa->update(); // 更新当前节点的size（仅自身cnt变化，无需调整平衡）
            return fa;    // 直接返回，不参与后续的更新/旋转（结构无变化）
        }
        
        fa->update(); // 回溯层：每层都执行，更新当前节点的size/height/x
        if(abs(fa->x) > 1) fa = fix(fa); // 平衡因子绝对值>1则失衡，旋转修复,更新树根 
        return fa; // 每层返回更新/修复后的当前节点（供上层更新子指针）
    }
    
    // 递归删除值为x的节点：核心逻辑=找节点+删节点+回溯修复平衡
    node* remove(node *cur, int x) {
        if (!cur) return nullptr; // 递归终止：没找到值为x的节点，直接返回空
        
        if (x < cur->d) cur->l = remove(cur->l, x); // 值更小，递归往左子树找并删除，返回新根更新左子节点 
        else if(x > cur->d) cur->r = remove(cur->r, x); // 值更大，往右子树找并删除
        else {// 找到要删除的节点
            if (cur->cnt > 1) {// 该值有重复，仅减少计数（不删除节点），无递归，仅一次 
                cur->cnt--;    // 重复数-1
                cur->update(); // 更新size（仅自身cnt变化）
                return cur;    // 返回自身，不删除节点，结构不变 
            }
            
            node *tmp = nullptr; // 用于替换当前节点的临时节点
            if (cur->l && cur->r) {// 双孩子节点：用后继节点替换（避免破坏BST结构）
                tmp = fsuc(cur); // 找后继节点（右子树的最左节点，比当前值大的最小值）
                cur->d = tmp->d;cur->cnt = tmp->cnt; // 后继替代当前节点
                tmp->cnt = 1; // 强制设为1：让后续删除时必物理删除（避免仅减计数）
                cur->r = remove(cur->r, tmp->d); // 递归删除后继节点（已无存在必要）
            }else{// 单孩子用子节点替换当前节点或，无孩子就直接删 ，无递归，仅一次 
                tmp = cur->l ? cur->l : cur->r; // 有左取左，有右取右，无则空
                node *fa = cur->f; // 取当前节点的父节点
                if (fa) {// 绕过当前节点，建立父节点和子节点的直接连接
                    if (fa->l == cur) fa->l = tmp;
                    else fa->r = tmp;
                } else root = tmp; // cur就是根，现在tmp替代成根
                if (tmp) tmp->f = fa; // 给子节点绑定新的父节点
                delete cur; cur = nullptr; // 物理删除当前节点，避免内存泄漏
                return tmp; // 返回替换节点，供上层更新子指针
            }
        }
        // 递归几次回溯几次：删除后更新平衡（仅当前节点未被删除时执行）
        if (cur) {
            cur->update(); // 更新size/height/x（子树结构变化）
            if (abs(cur->x) > 1) cur = fix(cur); // 失衡则旋转修复
        }
        return cur; // 返回更新/修复后的节点（供上层更新子指针）
    }
    
    // 右旋操作：fa节点下沉，其左子节点上浮为新根
    node* rright(node *fa){
        node *gfa = fa->f;   // fa的父节点（祖父节点）
        node *cur = fa->l;   // fa的左子节点（即将上浮的新根）
        node *gson = cur->r; // cur的右子节点（即将过继给fa）
        
        fa->l = gson;        // cur的右子过继给fa当左子
        if(gson) gson->f = fa; // 给过继的节点绑定父节点
        
        cur->f = gfa;        // 新根cur绑定原祖父节点
        if(gfa){// 祖父节点存在：更新祖父的子指针为新根
            if(gfa->l == fa) gfa->l = cur;
            else if(gfa->r == fa) gfa->r = cur;
        }else root = cur;    // 原fa是根节点：更新整棵树的根为cur
        
        fa->f = cur;         // 原fa绑定新根cur为父节点
        cur->r = fa;         // 原fa成为cur的右子节点
        
        fa->update();  // 先更新下沉节点的属性（子树结构变化）
        cur->update(); // 再更新新根的属性
        return cur;    // 返回旋转后的新根节点
    }
    
    // 左旋操作：fa节点下沉，其右子节点上浮为新根（逻辑与右旋对称）
    node* rleft(node *fa){
        node *gfa = fa->f;   // fa的父节点（祖父节点）
        node *cur = fa->r;   // fa的右子节点（即将上浮的新根）
        node *gson = cur->l; // cur的左子节点（即将过继给fa）
        
        fa->r = gson;        // cur的左子过继给fa当右子
        if(gson) gson->f = fa; // 给过继的节点绑定父节点
        
        cur->f = gfa;        // 新根cur绑定原祖父节点
        if(gfa){// 祖父节点存在：更新祖父的子指针为新根
            if(gfa->l == fa) gfa->l = cur;
            else if(gfa->r == fa) gfa->r = cur;
        }else root = cur;    // 原fa是根节点：更新整棵树的根为cur
        
        fa->f = cur;         // 原fa绑定新根cur为父节点
        cur->l = fa;         // 原fa成为cur的左子节点
        
        fa->update();  // 先更新下沉节点的属性
        cur->update(); // 再更新新根的属性
        return cur;    // 返回旋转后的新根节点
    }
    
    // 失衡修复：根据平衡因子判断旋转类型，返回修复后的子树根节点
    node* fix(node *fa){
        if(fa->x == 2){// 左子树过高（失衡类型：LL或LR）
            if(fa->l->x >= 0){// LL型失衡：fa直接右旋即可
                fa = rright(fa);//返回的是fa的子节点，更新fa 
            }else if(fa->l->x == -1){// LR型失衡：先左旋fa的左子，再右旋fa
                fa->l = rleft(fa->l);
                fa = rright(fa);
            }
        }else if(fa->x == -2){// 右子树过高（失衡类型：RR或RL）
            if(fa->r->x <= 0){// RR型失衡：fa直接左旋即可
                fa = rleft(fa);
            }else if(fa->r->x == 1){// RL型失衡：先右旋fa的右子，再左旋fa
                fa->r = rright(fa->r);
                fa = rleft(fa);
            }
        }
        return fa; //返回新根 
    }
    
    // 找后继节点：当前节点的右子树中最左的节点（比当前值大的最小值）
    node *fsuc(node *cur){
        if(!cur) return NULL;
        node *tmp = cur->r; // 先找右子树
        while(tmp->l) tmp = tmp->l; // 一直往左找，直到无左子（最小值）
        return tmp;
    }
    
    // 找值x的排名：排名=比x小的元素总数+1
    int get_rank(node *cur, int x){
        if (!cur) return 1; // 递归基：空树，x的排名为1（无元素比它小）
        int l_size = cur->l ? cur->l->size : 0; // 左子树的总元素数（全比当前值小）
        if(x < cur->d)return get_rank(cur->l, x);// x更小：递归到左子树找排名
        else if(cur->d < x)return l_size + cur->cnt + get_rank(cur->r, x);// x更大：排名=左子树总数+当前节点重复数+右子树中x的排名
        else return l_size + 1;// x等于当前值：排名=左子树总数+1
    }
    
    // 找第x名的元素：按中序遍历的顺序（从小到大）找第x个
    int get_xth(node *cur, int x) {
        if(!cur) return -1; // 递归基：空树/超出总数，返回-1（无效值）
        int l_size = cur->l ? cur->l->size : 0; // 左子树的总元素数
        if(x <= l_size)return get_xth(cur->l, x);// x在左子树范围内：递归左子树找
        else if(x <= l_size + cur->cnt) return cur->d;// x落在当前节点的重复数范围内：当前值就是第x名
        else return get_xth(cur->r, x - l_size - cur->cnt);// x在右子树范围内：递归右子树，且x需减去左子树+当前节点的总数
    }
    
    // 找x的前驱：比x小的最大值
    int get_pre(node *cur, int x)const{
        int res = INT_MIN; // 初始化为最小整数（无符合条件值时返回此值）
        while(cur){// 循环遍历，不递归（效率更高）
            if(cur->d < x){// 当前值更小：属于前驱候选
                res = max(res, cur->d); // 保留最大的候选值
                cur = cur->r; // 往右找更大的前驱（逼近x）
            }else  cur = cur->l;// 当前值>=x：往左找更小的值
        }
        return res;
    }
    
    // 找x的后继：比x大的最小值（逻辑与前驱对称）
    int get_suc(node *cur, int x)const{
        int res = INT_MAX; // 初始化为最大整数（无符合条件值时返回此值）
        while(cur){// 循环遍历，不递归
            if(cur->d > x){// 当前值更大：属于后继候选
                res = min(res, cur->d); // 保留最小的候选值
                cur = cur->l; // 往左找更小的后继（逼近x）
            }else cur = cur->r;// 当前值<=x：往右找更大的值
        }
        return res;
    }
    
    // 中序遍历打印树：按从小到大顺序输出节点信息
    void view(node *cur){
        if(!cur) return;
        if(cur->l) view(cur->l); // 先遍历左子树（更小值）
        // 打印当前节点：值、高度、总元素数、平衡因子
        cout << cur->d << ",高" << cur->height << ",数" << cur->size << ",因" << cur->x << "\t";
        if(cur->r) view(cur->r); // 再遍历右子树（更大值）
    }
    
    // 递归销毁整棵树：后序遍历（先删子树，再删根），避免内存泄漏
    void destroy(node *cur){
        if(!cur) return;
        if(cur->l) destroy(cur->l); // 先销毁左子树
        if(cur->r) destroy(cur->r); // 再销毁右子树
        delete(cur); cur = nullptr; // 最后销毁当前节点
    }

public:// 公共接口（外部可调用）
    avltree(){root = nullptr;} // 构造函数：初始化根节点为空
    ~avltree(){destroy(root);} // 析构函数：销毁整棵树，释放内存
    void insert(int x){root = insert(root, x);} // 插入接口：调用私有insert，更新根
    void remove(int x){root = remove(root, x);} // 删除接口：调用私有remove，更新根
    int get_rank(int x){return get_rank(root, x);} // 查排名接口
    int get_xth(int x){return get_xth(root, x);}   // 查第x名接口
    int get_pre(int x){return get_pre(root, x);}   // 查前驱接口
    int get_suc(int x){return get_suc(root, x);}   // 查后继接口
    
    // 打印整棵树的根和所有节点信息
    void view(){
        cout << "AVL"; 
        if(root) cout << root->d << endl; // 打印根节点的值
        if(!root) return;
        view(root); // 中序遍历打印所有节点
        cout << endl;
    } 
}tree;

int n, op, x;
int main(){
    //freopen("data.cpp","r",stdin); // 调试用：从文件读入数据
    cin >> n; // 操作总数
    while(n--){
        node *cur; // 临时变量（未使用，保留你的写法）
        cin >> op >> x; // 读入操作类型和操作数
        
        switch(op){
            case 1:
                tree.insert(x); // 操作1：插入值x
                break;
            case 2:
                tree.remove(x); // 操作2：删除值x
                break;
            case 3:
                cout << tree.get_rank(x) << endl;// 操作3：输出x的排名（比x小的元素数+1）
                break;
            case 4:
                // 操作4：输出第x小的元素（原注释写“第x大”是错误，已修正）
                cout << tree.get_xth(x) << endl;
                break;
            case 5:
                cout << tree.get_pre(x) << endl;// 操作5：输出x的前驱（比x小的最大值）
                break;
            case 6:
                // 操作6：输出x的后继（比x大的最小值）（原注释写“前驱”，已修正）
                cout << tree.get_suc(x) << endl;
                break;
        }
        //tree.view(); // 调试用：打印树的当前状态
    }
    return 0;
}