[平衡树] aw3786. 二叉排序树(BST)

1. 题目来源

链接：3786. 二叉排序树

C++版，好久以前写的了：[C++系列] 76. 详解BST二叉搜索树

2. 题目解析

链式二叉搜索树的简单实现，也是非常经典的数据结构。

支持常见三种操作：增加、删除、查找。树中元素不重复，若有重复值可通过打点标记计数，将重复值归到左子树或右子树等操作完成。

思路：

• 增加：只会在叶子节点增加新节点，类比二分查找思路，递归操作即可。
• 删除：三种情况，操作保证中序遍历仍有序均可。
  • 叶子节点：直接删除即可
  • 仅有左子树、仅有右子树：左右子树替代当前节点即可。
  • 左右子树均存在，找到左子树中最大值，即左孩子最右链中的节点，将其值代替，然后删除该节点即可。值替换，节点替代删除。
    • 注意，在删除过程中，不能直接将替代节点赋为 NULL 代表删除，因为这个节点无右孩子但仍可能存在左孩子。应该递归删除左子树中值为替代值的这个节点，由于值唯一，一定能够找到，此时一定不存在右孩子，可能存在左孩子，左孩子构成一棵树，这是一个递归定义。
    • 当然，迭代写法也是完全可以的，用 prev 指针指向替换节点的父节点，若删除替换节点，等价于将 prev 的右指针指向替换节点的左子树。有边界情况需要考虑，若 prev==root 即，root 的左孩子没有右子树时，应该删除的是左孩子，则，应该为 root->left = p->left，此时 prev 为空。
• 查找：两种情况，查找前驱节点、查找后继节点。前驱节点：predecessor，后继节点：successor
  • 前驱节点：找到小于 x 的最大值。类比二分查找即可，如果根节点值大于等于 x ，那么应该进入其左子树查找，左子树的所有值都严格小于 根节点值，可能会小于 x。否则，根节点的值已经严格小于 x，答案可能是根节点，也可能是根节点的右子树中最大的那个节点值。左子树由于更进一步严格小于 x 故不再考虑范围内。当找到一个空节点时，我们返回一个极小值不影响答案即可，因为一个节点的右子树可能是不存在的，取 max 的时候不受影响即可。
  • 后继节点：找到大于 x 的最小值。道理和前驱节点一样，如果根节点值小于等于 x 则答案在右子树中。否则，答案可能是根节点和根节点的左子树中的最小值。找不到的时候返回极大值，防止取 min 的时候出现干扰即可。

---

BST 平衡树，理解其含义即可。代码可以多码码，毕竟是最基础的平衡树，还没上旋转呢hh。

---

时间复杂度：$O(h)$

空间复杂度：$O(n)$

---

#include <bits/stdc++.h>

using namespace std;

const int INF = 1e9;

struct TreeNode {
    int val;
    TreeNode *left, *right;
    TreeNode(int _val) : val(_val), left(NULL), right(NULL) {}
} *root;

// 本题数值保证各不相同
void insert(TreeNode *&root, int x) {
    if (!root) root = new TreeNode(x);
    else if (root->val > x) insert(root->left, x);
    else insert(root->right, x);
}

void erase(TreeNode *&root, int x) {
    if (!root) return ;
    else if (root->val > x) erase(root->left, x);
    else if (root->val < x) erase(root->right, x);
    else {
        if (!root->left && !root->right) root = NULL;       // 叶节点直接删除
        else if (!root->left) root = root->right;           // 单右儿子，右儿子当树
        else if (!root->right) root = root->left;           // 单左儿子，左儿子当树
        else {                                              // 左右儿子都存在，找左子树最大值，替换删除
            TreeNode *p = root->left;
            
            /* 迭代写法
            TreeNode *prev = NULL;
            while (p->right) prev = p, p = p->right;
            root->val = p->val;
            if (prev) prev->right = p->left;
            else root->left = p->left;
            */
            
            // 递归简洁明了
            erase(root->left, p->val);          // 在这不要直接 p=NULL, 因为 p 可能还存有左子树，并不一定是叶子节点
        }
    }
}

// 找前驱节点，predecessor
int get_pre(TreeNode *root, int x) {
    if (!root) return -INF;                                 // 未找到合法前驱
    if (root->val >= x) return get_pre(root->left, x);
    
    return max(root->val, get_pre(root->right, x));
}

// 找后继节点，successor
int get_suc(TreeNode *root, int x) {
    if (!root) return INF;
    if (root->val <= x) return get_suc(root->right, x);
    
    return min(root->val, get_suc(root->left, x));
}

int main() {
    int T; cin >> T; while (T -- ) {
        int op, val;
        cin >> op >> val;
        if (op == 1) insert(root, val);
        else if (op == 2) erase(root, val);
        else if (op == 3) cout << get_pre(root, val) << endl;
        else cout << get_suc(root, val) << endl;
    }
    
    return 0;
}