[AcWing]3786. 二叉排序树

你需要写一种数据结构，来维护一些数，其中需要提供以下操作：

  1. 插入数值 x。
  2. 删除数值 x。
  3. 输出数值 x 的前驱(前驱定义为现有所有数中小于 x 的最大的数)。
  4. 输出数值 x 的后继(后继定义为现有所有数中大于 x 的最小的数)。

题目保证：

  • 操作 1 插入的数值各不相同。
  • 操作 2 删除的数值一定存在。
  • 操作 3 和 4 的结果一定存在。

输入格式

第一行包含整数 n，表示共有 n 个操作命令。

接下来 n 行，每行包含两个整数 opt 和 x，表示操作序号和操作数值。

输出格式

对于操作 3,4，每行输出一个操作结果。

数据范围

1 < n < 2000，
-10000 < x < 10000

输入样例：

1
2
3
4
5
6
7
6
1 1
1 3
1 5
3 4
2 3
4 2

输出样例：

1
2
3
5

算法思想

略. 详细看代码注释

代码实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
#include <iostream>

using namespace std;

const int INF = 0x3f3f3f3f;

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

// 要用引用，需要将改变的结果带回去
void insertNode(TreeNode *& root, int x){
    // 此题没有重复的元素插入
    // 遇到空节点直接创建
    if(!root) root = new TreeNode(x);
    else if(x < root -> val) insertNode(root -> left, x); // 如果插入的值比当前节点的值小，就去左子树找
    else insertNode(root -> right, x); // 否则去右子树
}

void deleteNode(TreeNode *& root, int x){
    // 都找遍了都没找到这个结点，那么干脆什么都不做了
    if(!root) return;
    // 如果要删除的节点比当前节点小，去左子树找，否则去右子树找
    else if(x < root -> val) deleteNode(root -> left, x);
    else if(x > root -> val) deleteNode(root -> right, x);
    else{// 找到了
        // 1. 刚好是叶子节点，直接删除
        if(!root -> left && !root -> right) delete root, root = NULL;
        // 2. 只有左子树，只有右子树。让子树接替
        else if(!root -> right) root = root -> left;
        else if(!root -> left) root = root -> right;
        else{// 3. 左右子树都存在
            // root的前驱节点顶替root
            TreeNode * p = root -> left;
            while(p -> right) p = p -> right;
            root -> val = p -> val;
            // 去左子树当中删掉它的前驱
            // 在你找到 p 后，简单地删除 p 会破坏树的结构，因为它可能还有子节点。
            deleteNode(root -> left, p -> val);
        }
    }
}

// 找前驱（小于x的最大数）
// 当前值，比x大，应该往左递归，继续找
// 当前值，比x小，当前值可能作为答案，当前节点的右子树中也可能存在答案，两者取max即可
int findPre(TreeNode * root, int x){
    // INF是不可能取的值，代表没有找到前驱
    if(!root) return -INF;
    if(root -> val >= x) return findPre(root -> left, x);
    return max(root -> val, findPre(root -> right, x));
}

// 找后继（大于x的最小数）
// 当前值，比x小，应该往右递归，继续找
// 当前值，比x大，当前值可能作为答案，当前节点的左子树中也可能存在答案，两者取min即可
int findSuc(TreeNode * root, int x){
    // INF是不可能取的值，代表没有找到后继
    if(!root) return INF;
    if(root -> val <= x) return findSuc(root -> right, x);
    return min(root -> val, findSuc(root -> left, x));
}

int main(){
    
    int n;
    cin >> n;
    
    TreeNode * root;  
    
    while(n --){
        int op, x; 
        cin >> op >> x;
        
        if(op == 1){
            // 插入
            insertNode(root, x);
        }else if(op == 2){
            // 删除
            deleteNode(root, x);
        }else if(op == 3){
            // 前驱
            cout << findPre(root, x) << endl;
        }else if(op == 4){
            // 后继
            cout << findSuc(root, x) << endl;
        }
    }
        
    return 0;
}