算法数据结构

手撕平衡树-二叉查找树BST

字数统计: 2.7k字   |   阅读时长: 12分
, ,

摘要: 二叉查找树(BST)的原理、代码模板、例题

【对数据分析、人工智能、金融科技、风控服务感兴趣的同学，欢迎关注我哈，阅读更多原创文章】
我的网站：潮汐朝夕的生活实验室
我的公众号：潮汐朝夕
我的知乎：潮汐朝夕
我的github：FennelDumplings
我的leetcode：FennelDumplings

$1 二叉查找树定义

二叉查找树(Binary Search Tree)是指一棵空树或者具有下列性质的二叉树：

  • 对任意节点，若左子树不空，则左子树上所有节点的值均小于它的根节点的值；
  • 对任意节点，若右子树不空，则右子树上所有节点的值均大于它的根节点的值；
  • 任意节点的左、右子树也分别为二叉查找树；
  • 没有键值相等的节点。

中序遍历二叉查找树可得到一个关键字的有序序列，一个无序序列可以通过构造一棵二叉查找树变成一个有序序列，构造树的过程即为对无序序列进行查找的过程。

二叉查找树可以用于构建更抽象的数据结构，如集合、多重集合、关联数组等。

二叉查找树的高度 $h_{N} : 由 N 个节点随即构成的二叉查找树的高度$，以下为结论

  1. 树高 $O(\log N)$ 是概率性的
  2. 一系列插入删除后会改变树的结构，这种改变，最坏情况会使树退化为链表

为了避免退化，有很多改进版的二叉查找树，通过在节点上加字段以及相应的操作，使得插入删除时候可以维持树的平衡，称为平衡树。主流的有 AVL 树和红黑树等。

$2 二叉搜索树的实现

节点定义：

1
2
3
4
5
6
7
8
9
10
11
12
13
struct BSTNode
{
    int data;
    BSTNode *left;
    BSTNode *right;
    BSTNode():left(nullptr),right(nullptr){}
    BSTNode(const int& x, BSTNode* p=nullptr, BSTNode* q=nullptr):data(x),left(p),right(q){}
    ~BSTNode()
    {
        left = nullptr;
        right = nullptr;
    }
};

二叉查找树三个基础功能，插入、删除、查找的接口。

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
class BinarySearchTree
{
public:
    BinarySearchTree():root(nullptr){}
    ~BinarySearchTree()
    {
        if(root)
            _delete_sub_tree(root);
    }

    bool find(const int& x) const
    {
        return find(x, root);
    }

    void insert(const int& x)
    {
        insert(x, root);
    }

    void remove(const int& x)
    {
        remove(x, root);
    }

private:
    BSTNode *root;

    void _delete_sub_tree(BSTNode* node)
    {
        if(node -> left)
            _delete_sub_tree(node -> left);
        if(node -> right)
            _delete_sub_tree(node -> right);
        delete node;
        node = nullptr;
    }

    void insert(const int& x, BSTNode*& t);
    void remove(const int& x, BSTNode*& t);
    BSTNode* find(const int& x, BSTNode* t) const;
};

以下为二叉查找树插入，删除，查找的实现，力扣上有对应的题目：

查找

1
2
3
4
5
6
7
BSTNode* BinarySearchTree::find(const int& x, BSTNode* t) const
{
    if(t == nullptr) return nullptr;
    else if(t -> data > x) return find(x, t -> left);
    else if(t -> data < x) return find(x, t -> right);
    else return t;
}

插入

在递归中，父节点的连接有参数传递完成：

1
2
3
4
5
6
7
8
9
void BinarySearchTree::insert(const int& x, BSTNode*& t)
{
    if(t == nullptr)
        t = new BSTNode(x, nullptr, nullptr);
    else if(t -> data > x)
        insert(x, t -> left);
    else
        insert(x, t -> right);
}

删除

首先找到被删节点 t，若 t 无子节点，或只有一个子节点，可以直接将 t 赋值为 t 的子节点，再将原来的 t 表示的节点删掉：

1
2
3
BSTNode *oldNode = t;
t = (t -> left != nullptr) ? t -> left : t -> right;
delete oldNode;

若 t 有两个子节点，先找到 t 的前驱或后继：

1
2
3
4
5
6
7
8
// 后继
BSTNode *successor = t -> right;
while(successor -> left != nullptr)
    successor = successor -> left;
// 前驱
TreeNode* precursor = t -> left;
while(precursor -> right)
    precursor = precursor -> right;

将后继的节点值给 t，删除右子树中后继的节点值：

1
2
t -> data = successor -> data;
remove(t -> data, t -> right);

或者将前驱的节点值给 t，删除左子树子树中前驱的节点值：

1
2
t -> data = precursor -> data;
remove(t -> data, t -> left);

删除值为 x 的节点的完整代码：

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
void BinarySearchTree::remove(const int& x, BSTNode*& t)
{
    if(t == nullptr) 
        return;
    if(x < t -> data)
        remove(x, t -> left);
    if(x > t -> data)
        remove(x, t -> right);
    if(x == t -> data)
    {
        // 找到被删节点 t
        if(t -> left != nullptr && t -> right != nullptr)
        {
            // t 有两个子节点
            BSTNode *successor = t -> right;
            while(successor -> left != nullptr)
                successor = successor -> left;
            t -> data = successor -> data;
            remove(t -> data, t -> right);
        }
        else
        {
            // t 只有 1 个子节点，或没有子节点
            BSTNode *oldNode = t;
            t = (t -> left != nullptr) ? t -> left : t -> right;
            delete oldNode;
        }
    }
}

$3 带 size 字段的二叉查找树

$2 的二叉查找树只是基本数据结构，它可以解决基本的插入，删除，查找需求。

但是有一些 BST 上有意义的需求，则需要在节点中加上某些字段才能使统计过程更高效。

例如：问 BST 中小于于 x 的元素个数有多少，或者问 BST 中第 k 小的元素。

这两个问题，如果树节点有一个 size 字段，表示该节点对应子树的元素个数，则 $O(\log N)$ 就可以完成统计。而如果是基础的二叉查找树，需要在中序遍历过程中每经过一个点做一次统计，需要 $O(N)$ 完成统计。

这两个需求在力扣上有对应题目：

size 字段本身可以用于 BST 的平衡操作，使得 BST 成为平衡树，这种基于 size 字段的平衡树称为大小平衡树 SBT(Size Balanced Tree)。

类似地，有的与树高度相关的需求，需要给树节点加上高度 height 字段可以更高效地解决，而 height 字段本身也可以用于 BST 的平衡操作，使得 BST 成为平衡树。这种基于 height 字段的平衡树称为 AVL 树。

下面实现带 size 字段的二叉查找树并用于解力扣第 315 题。

带 size 字段的 BST

  • 节点定义：除了增加 size 字段，没有变化
1
2
3
4
5
6
7
8
9
10
11
12
13
14
struct BSTNode
{
    int data;
    BSTNode *left;
    BSTNode *right;
    int size;
    BSTNode():left(nullptr),right(nullptr){}
    BSTNode(const int& x, BSTNode* p=nullptr, BSTNode* q=nullptr, int s=1):data(x),left(p),right(q),size(s){}
    ~BSTNode()
    {
        left = nullptr;
        right = nullptr;
    }
};

  • 查找不变。

  • 插入和删除时需要动态维护 size，注意 ++(t -> size) 和 --(t -> size) 的位置;

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
void BinarySearchTree::insert(const int& x, BSTNode*& t)
{
    if(t == nullptr)
    {
        t = new BSTNode(x, nullptr, nullptr, 1);
        return;
    }
    ++(t -> size);
    if(t -> data > x)
        insert(x, t -> left);
    else
        insert(x, t -> right);
}

void BinarySearchTree::remove(const int& x, BSTNode*& t)
{
    if(t == nullptr) return;
    --(t -> size);
    if(x < t -> data)
        remove(x, t -> left);
    if(x > t -> data)
        remove(x, t -> right);
    if(x == t -> data)
    {
        // 找到被删节点 t
        if(t -> left != nullptr && t -> right != nullptr)
        {
            // t 有两个子节点
            BSTNode *successor = t -> right;
            while(successor -> left != nullptr)
                successor = successor -> left;
            t -> data = successor -> data;
            remove(t -> data, t -> right);
        }
        else
        {
            // t 只有 1 个子节点，或没有子节点
            BSTNode *oldNode = t;
            t = (t -> left != nullptr) ? t -> left : t -> right;
            delete oldNode;
        }
    }
}
  • 需求查询：小于 x 的元素有多少个。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
int BinarySearchTree::lessthan(const int& x, BSTNode* t) const
{
    if(t -> data >= x)
    {
        if(!t -> left)
            return 0;
        return lessthan(x, t -> left);
    }
    // t -> data < x
    int ans = 1;
    if(t -> left)
        ans += t -> left -> size;
    if(t -> right)
        ans += lessthan(x, t -> right);
    return ans;
}

以上是带 size 字段的 BST，与基本 BST 的区别是节点定义增加 size 字段，插入删除时需要动态维护每个节点的 size 字段。此外增加了大于 x 的元素个数的查询接口。

以下是利用这颗带 size 字段的 BST 解力扣315 题，315. 计算右侧小于当前元素的个数，本题只用到了插入，和需求功能的查询。

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
94
95
96
97
98
99
100
struct BSTNode
{
    int data;
    BSTNode *left;
    BSTNode *right;
    int size;
    BSTNode():left(nullptr),right(nullptr){}
    BSTNode(const int& x, BSTNode* p=nullptr, BSTNode* q=nullptr, int s=0):data(x),left(p),right(q),size(s){}
    ~BSTNode()
    {
        left = nullptr;
        right = nullptr;
    }
};

class BinarySearchTree
{
public:
    BinarySearchTree():root(nullptr){}
    ~BinarySearchTree()
    {
        if(root)
            _delete_sub_tree(root);
    }

    void insert(const int& x)
    {
        insert(x, root);
    }

    int lessthan(const int& x) const
    {
        return lessthan(x, root);
    }

private:
    BSTNode *root;

    void _delete_sub_tree(BSTNode* node)
    {
        if(node -> left)
            _delete_sub_tree(node -> left);
        if(node -> right)
            _delete_sub_tree(node -> right);
        delete node;
        node = nullptr;
    }

    int lessthan(const int& x, BSTNode* t) const;
    void insert(const int& x, BSTNode*& t);
};

int BinarySearchTree::lessthan(const int& x, BSTNode* t) const
{
    if(t -> data >= x)
    {
        if(!t -> left)
            return 0;
        return lessthan(x, t -> left);
    }
    // t -> data < x
    int ans = 1;
    if(t -> left)
        ans += t -> left -> size;
    if(t -> right)
        ans += lessthan(x, t -> right);
    return ans;
}

void BinarySearchTree::insert(const int& x, BSTNode*& t)
{
    if(t == nullptr)
    {
        t = new BSTNode(x, nullptr, nullptr, 1);
        return;
    }
    ++(t -> size);
    if(t -> data > x)
        insert(x, t -> left);
    else
        insert(x, t -> right);
}

class Solution {
public:
    vector<int> countSmaller(vector<int>& nums) {
        if(nums.empty()) return {};
        int n = nums.size();
        vector<int> result(n);
        BinarySearchTree bst;
        result[n - 1] = 0;
        bst.insert(nums[n - 1]);
        for(int i = n - 2; i >= 0; --i)
        {
            result[i] = bst.lessthan(nums[i]);
            bst.insert(nums[i]);
        }
        return result;
    }
};

$4 二叉查找树题目汇总

二叉查找树的中序遍历和前驱后继

树的DFS遍历

$5 二叉查找树节点的旋转

为了保持二叉搜索树的平衡,我们通常通过旋转改变指针结构。这种旋转是一种可以保持二叉搜索树特性的本地运算。

  • 左旋 Left-Ratote(x) 操作通过更改两个指针将左边两个结点的结构转变成右边的结构,
  • 右边的结构也可以通过相反的操作 Right-Ratote(x) 来转变成左边的结构。

右旋 - 需要左子节点存在

1
2
3
4
5
6
7
8
void right_rotate(BSTNode*& t)
{
    // 调用方保证 t -> left 存在
    BSTNode *tmp = t -> left;
    t -> left = tmp -> right;
    tmp -> right = t;
    t = tmp;
}

左旋 - 需要右子节点存在

1
2
3
4
5
6
7
8
void left_rotate(BSTNode*& t)
{
    // 调用方保证 t -> right 存在
    BSTNode *tmp = t -> right;
    t -> right = tmp -> left;
    tmp -> left = t;
    t = tmp;
}

带 size 字段的左旋和右旋

左旋和右旋的改动均是在 t = tmp 之前加入以下两步对 size 的更新：

1
2
3
4
5
t -> size = 1;
if(t -> right)
    t -> size += t -> right -> size;
if(t -> left)
    t -> size += t -> left -> size;
1
2
3
4
5
6
7
8
9
10
11
12
13
14
void right_rotate(BSTNode*& t)
{
    // 调用方保证 t -> left 存在
    BSTNode *tmp = t -> left;
    t -> left = tmp -> right;
    tmp -> right = t;
    tmp -> size = t -> size;
    t -> size = 1;
    if(t -> right)
        t -> size += t -> right -> size;
    if(t -> left)
        t -> size += t -> left -> size;
    t = tmp;
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
void left_rotate(BSTNode*& t)
{
    // 调用方保证 t -> right 存在
    BSTNode *tmp = t -> right;
    t -> right = tmp -> left;
    tmp -> left = t;
    tmp -> size = t -> size;
    t -> size = 1;
    if(t -> right)
        t -> size += t -> right -> size;
    if(t -> left)
        t -> size += t -> left -> size;
    t = tmp;
}
Share