shan

二叉搜索树及平衡二叉搜索树

2020-06-06

本文为 python数据结构与算法的读书笔记，主要内容是二叉搜索树及平衡二叉搜索树的实现方式。

二叉搜索树

这是有种从集合中获取键值对的方法。这里不关心元素在树中的确切位置，而是如何利用二叉树结构提供高效的搜索。

搜索树的操作

  • Map() 新建一个空的映射
  • put(key, val) 往映射中加入一个新的键值对。如果键已经存在了，就用新值替换旧值
  • get(key) 返回key对应的值，如果key不存在，返回None。
  • del 通过del map[key] 这样的语句从映射中删除键值对
  • len() 返回映射中存储的键值对的数目
  • in 通过key in map 这样的语句，在键存在时，返回True，否则返回False

搜索树的实现

二叉搜索树的性质：小于父节点的键都在左子树中，大于父节点的键都在右子树中，我们称之为二叉搜索性，它会引导我们实现上述映射接口。下面使用节点与引用实现二叉搜索树。

我们使用两个类，一个称为BinarySearchTree，另一个称为TreeNode。BinarySearchTree类有一个引用，指向作为二叉搜索树根节点的TreeNode类。大多数情况下，外面这个类的方法只检查树是否为空。如果树中有节点，请求就会被发往BinarySearchTree类的私有方法，这个方法以根节点作为参数。当树为空，或者想删除根节点的键时，需要采取特殊措施。

实现TreeNode类

下面是TreeNode类，它提供了大量的辅助方法，简化了BinarySearchTree类的工作。下面是相关的代码：

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
101
102
103
104
105
106
107
108
109

class TreeNode:
    def __init__(self, key, val, left=None, right=None, parent=None):
        self.key = key
        self.payload = val
        self.left_child = left
        self.right_child = right
        self.parent = parent

    def has_left_child(self):
        """
        是否有左节点
        :return:
        """
        return self.left_child

    def has_right_child(self):
        """
        是否有右节点
        :return:
        """
        return self.right_child

    def is_left_child(self):
        """
        当前节点是否为左节点
        :return:
        """
        return self.parent and self.parent.left_child == self

    def is_right_child(self):
        """
        当前节点是否为右节点
        :return:
        """
        return self.parent and self.parent.right_child == self

    def is_root(self):
        """
        是否为根节点
        :return:
        """
        return not self.parent

    def is_leaf(self):
        """
        是否为叶子节点
        :return:
        """
        return not(self.right_child or self.left_child)

    def has_any_children(self):
        """
        是否存在左或右节点
        :return: 
        """
        return self.left_child or self.right_child

    def has_both_children(self):
        """
        是否存在左和右节点
        :return: 
        """
        return self.left_child and self.right_child

    def replace_node_data(self, key, value, lc, rc):
        """
        替换节点数据
        :param key: 需替换节点的键值
        :param value: 需替换节点的value值
        :param lc: 左节点
        :param rc: 右节点
        :return: 
        """
        self.key = key
        self.payload = value
        self.left_child = lc
        self.right_child = rc
        if self.has_left_child():
            self.left_child.parent = self
        if self.has_right_child():
            self.right_child.parent = self
        
	def find_successor(self):
        """
        寻找后继节点
        """
        succ = None
        if self.has_right_child():
            succ = self.right_child.findMin()
        else:
            if self.parent:
                if self.is_left_child():
                    succ = self.parent
                else:
                    self.parent.rightChild = None
                    succ = self.parent.findSuccessor()
                    self.parent.rightChild = self
        return succ

    def find_min(self):
        """
        寻找最小节点
        """
        current = self
        while current.hasLeftChild():
            current = current.leftChild
    
        return current

BinarySearchTree 基本结构

下面是主要搜索树BinarySearchTree的基本结构，之后会添加各种相关操作的方法：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
class BinarySearchTree:

    def __init__(self):
        self.root = None
        self.size = 0

    def length(self):
        return self.size

    def __len__(self):
        return self.size

    def __iter__(self):
        return self.root.__iter__()

put

首先添加一个put方法，它检查树是否已经存在根节点，若没有，则将其视为树的根节点，若有，就调用私有的递归辅助方法_put，并根据下面的算法在树中搜索：

  • 从根节点开始搜索二叉树，比较新键与当前节点的键，如果新键更小，搜索左子树，如果新键更大，搜索右子树。
  • 当没有可供搜索的左子节点时，就说明找到了新键的插入位置。
  • 向树中插入一个节点，做法是创建一个TreeNode对象，并将其插入到前一步发现的位置上。

也可以在类中添加__setitem__() 方法，来重载python的[]操作符，其相关的方法代码如下：

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
class BinarySearchTree:
    """
    这里省略了已有的代码
    """
    def __setitem__(self, key, value):
        self.put(key, value)

    def put(self, key, value):
        if self.root:
            self._put(key, value, self.root)
        else:
            self.root = TreeNode(key, value)
        self.size += 1

    def _put(self, key, value, current_node):
        if key < current_node.key:
            if current_node.has_left_child():
                self._put(key, value, current_node.left_child)
            else:
                current_node.left_child = TreeNode(key, value, parent=current_node)

        elif key > current_node.key:
            if current_node.has_right_child():
                self._put(key, value, current_node.right_child)
            else:
                current_node.right_child = TreeNode(key, value, parent=current_node)

        else:
            current_node.replace_node_data(key, value)

get

下一个任务就是为给定的键值去value，即实现get相关的方法，这里只需要递归的搜索二叉树，直到访问叶子节点或者找到匹配的键，然后返回对应的value。同样这里使用_get辅助函数，与_put 类似，但这里需要返回一个TreeNode。

并且可以使用__getitem__() 方法重载[]操作符，其相关代码如下：

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
class BinarySearchTree:
    """
    这里省略了已有的代码
    """
    def __getitem__(self, item):
        return self.get(item)
    
    def get(self, key):
        if self.root:
            node = self._get(key, self.root)
            if node:
                return node.payload
            else:
                return None
        else:
            return None

    def _get(self, key, current_node):
        if not current_node:
            return None
        elif current_node.key == key:
            return current_node
        elif key < current_node.key:
            self._get(key, current_node.left_child)
        elif key > current_node.key:
            self._get(key, current_node.right_child)
        else:
            return current_node

contains

我们也需要重写__contains__() 方法，来重载in运算符，以便于可以写出key in Tree这样的语句，其实现代码如下：

1
2
3
4
5
6
7
8
9
10
class BinarySearchTree:
    """
    这里省略了已有的代码
    """
    
    def __contains__(self, key):
        if self._get(key, self.root):
            return True
        else:
            return False

最后我们需要实现一个删除键的功能，首先它需要从树中搜索并找到要删除的节点。如果树中不止一个节点，使用_get() 方法搜索，找到需要移出的节点，如果找不到键，则delete方法会抛出异常。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
class BinarySearchTree:
   """
   这里省略了已有的代码
   """

   def delete(self, key):
       if self.size > 1:
           node_to_remove = self._get(key, self.root)
           if node_to_remove:
               self.remove(node_to_remove)
           else:
               raise KeyError("Error, key not in tree")
       elif self.size == 1 and self.root.key == key:
           self.root = None
           self.size = 0

       else:
           raise KeyError("Error, key not in tree")

remove

这里移出节点使用一个remove的辅助方法，具体分为3种情况：

  • 待删除节点没有子节点；
  • 待删除节点只有一个子节点；
  • 待删除节点有两个子节点。

具体细节如下：

  1. 待删除节点没有子节点，如果没有子节点，则直接删除这个节点，并移出父节点对这个节点的引用。
1
2
3
4
5
6
7
8
9
10
class BinarySearchTree:
    """
    这里省略了已有的代码
    """
    def remove(self, current_node):
        if current_node.is_leaf():
            if current_node == current_node.parent.left_child:
                current_node.parent.left_child = None
            else:
                current_node.parent.right_child = None
  1. 待删除节点只有一个子节点：
  • 当前节点是一个左子节点，只需要将当前节点的左子节点对父节点的引用改为指向当前节点的父节点，然后将父节点对当前节点的引用指向当前节点的左子节点。
  • 如果当前节点是一个右子节点，只需将当前节点的右子节点对父节点的引用该为指向单签节点的父节点，然后将父节点对当前节点的引用改为指向当前节点的右子节点。
  • 如果当前节点没有父节点，那么它是根节点，调用replace_node_data方法，替换根节点的key，payload，left_child，right_child数据。
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
class BinarySearchTree:
    """
    这里省略了已有的代码
    """
    
    def remove(self, current_node):
        if current_node.is_leaf():
            if current_node == current_node.parent.left_child:
                current_node.parent.left_child = None
            else:
                current_node.parent.right_child = None
        else:
            if current_node.has_left_child():
                if current_node.is_left_child():
                    current_node.left_child.parent = current_node.parent
                    current_node.parent.left_child = current_node.left_child
                elif current_node.is_right_child():
                    current_node.left_child.parent = current_node.parent
                    current_node.parent.right_child = current_node.left_child
                else:
                    current_node.replace_node_data(current_node.left_child.key, 
                                                   current_node.left_child.payload,
                                                   current_node.left_child.left_child,
                                                   current_node.left_child.right_child)
            else:
                if current_node.is_right_child():
                    current_node.right_child.parent = current_node.parent
                    current_node.parent.left_child = current_node.right_child
                elif current_node.is_left_child():
                    current_node.right_child.parent = current_node.parent
                    current_node.parent.right_child = current_node.right_child
                else:
                    current_node.replace_node_data(current_node.right_child.key,
                                                   current_node.right_child.payload,
                                                   current_node.right_child.left_child,
                                                   current_node.right_child.right_child)
  1. 待删除节点有两个子节点：这里可以搜索整棵树，找到可以替换待删除节点的节点，候选节点要能为左右子树都保持二叉搜索树的关系，也就是树中具有次大键的节点，我们称这个节点为后继节点。后继节点必定不会多于一个，所以我们知道如何按照已经实现的两种方法来删除它，移出后继节点后，只需要将它放到树中待删除节点的位置上即可。
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:
    """
    这里省略了已有的代码
    """
    
    def remove(self, current_node):
        if current_node.is_leaf():
            if current_node == current_node.parent.left_child:
                current_node.parent.left_child = None
            else:
                current_node.parent.right_child = None
		elif current_node.has_both_children():
            succ = current_node.find_duccessor()
            succ.splice_out()
            current_node.key = succ.key
            current_node.payload = succ.payload
            
        else:
            if current_node.has_left_child():
                if current_node.is_left_child():
                    current_node.left_child.parent = current_node.parent
                    current_node.parent.left_child = current_node.left_child
                elif current_node.is_right_child():
                    current_node.left_child.parent = current_node.parent
                    current_node.parent.right_child = current_node.left_child
                else:
                    current_node.replace_node_data(current_node.left_child.key, 
                                                   current_node.left_child.payload,
                                                   current_node.left_child.left_child,
                                                   current_node.left_child.right_child)
            else:
                if current_node.is_right_child():
                    current_node.right_child.parent = current_node.parent
                    current_node.parent.left_child = current_node.right_child
                elif current_node.is_left_child():
                    current_node.right_child.parent = current_node.parent
                    current_node.parent.right_child = current_node.right_child
                else:
                    current_node.replace_node_data(current_node.right_child.key,
                                                   current_node.right_child.payload,
                                                   current_node.right_child.left_child,
                                                   current_node.right_child.right_child)

这里需要对TreeNode添加find_successor，find_min来查找后继节点，对BinarySearchTree添加splice_out等辅助函数来移除后继节点，具体方法见TreeNode代码。寻找后继节点需要考虑3种情况：

  1. 如果节点有右子节点，那么后继节点就是右子树中最小的节点；
  2. 如果节点没有右子节点，并且其本身是父节点的左子节点，那么后继节点就是父节点；
  3. 如果节点是父节点的右子节点，并且其本身没有右子节点，那么后继节点就是除本身外父节点的后继节点。
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
class TreeNode:
	"""
    这里省略了已有的代码
    """
    
    def find_successor(self):
        succ = None
        if self.has_right_child():
            succ = self.right_child.find_min()
        else:
            if self.parent:
                if self.is_left_child():
                    succ = self.parent
                else:
                    self.parent.right_child = None
                    succ = self.parent.find_successor()
                    self.parent.right_child = self
        return succ

    def find_min(self):
        current = self
        while current.has_left_child():
            current = current.left_child

    def splice_out(self):
        if self.is_leaf():
            if self.is_left_child():
                self.parent.left_child = None
            else:
                self.parent.right_child = None
        elif self.has_any_children():
            if self.has_left_child():
                self.parent.left_child = self.left_child
            else:
                self.parent.right_child = self.left_child
        else:
            if self.is_left_child():
                self.parent.left_child = self.right_child
            else:
                self.parent.right_child = self.right_child
            self.right_child.parent = self.parent

为了支持迭代器操作，重写__iter__()方法，重载了for x in 操作，由于递归是发生在TreeNode上的，所以相关的方法也被定义在了这个类中：

1
2
3
4
5
6
7
8
9
10
11
12
13
class TreeNode:
	"""
    这里省略了已有的代码
    """
    def __iter__(self):
        if self:
            if self.has_left_child():
                for elem in self.left_child:
                    yield elem
            yield self.key
            if self.has_right_child():
                for elem in self.right_child:
                    yield elem

平衡二叉搜索树

这里介绍一种能够自动保持平衡的树，称为AVL树。其实现映射抽象数据类型的方式与普通的二叉搜索树一样，唯一的差别就是性能，实现AVL树，要记录每个节点的平衡因子。通过查看每个节点左右子树的高度来实现这一点，准确的说，我们将平衡因子定义为左右子树的高度差，其平衡因子大于0，称为左倾，小于0，称为右倾，等于0，那么树就是完全平衡的。为了实现AVL树并利用平衡树的优势，将平衡因子为-1，0，1的树都定义为平衡树。通过一系列的计算可以得知，AVL树的高度为节点数n取对数后乘以1.44。

AVL树的实现

AVL树与普通的二叉搜索树实现不同主要是插入和删除，因为AVL树需要记录平衡因子，并通过左旋或者右旋来维持树的平衡。这里我们将AVL树实现为二叉搜索树的子类(这里需要重写__init__()方法，加入平衡因子)，重写_put方法，并用新的update_balance方法来保持树的平衡，其代码如下：

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
class AVLTree(BinarySearchTree):
    def _put(self, key, value, current_node):
        if key < current_node.key:
            if current_node.has_left_child():
                self._put(key, value, current_node.left_child)
            else:
                current_node.left_child = TreeNode(key, value, parent=current_node)
                self.update_balance(current_node.left_child)
        else:
            if current_node.has_right_child():
                self._put(key, value, current_node.right_child)
            else:
                current_node.right_child = TreeNode(key, value, parent=current_node)
                self.update_balance(current_node.right_child)
    
    def update_balance(self, node):
        if node.balance_factor > 1 or node.balance_factor < -1:
            self.rebalance(node)
            return 
        if node.parent is not None:
            if node.is_left_child():
                node.parent.balance_factor += 1
            elif node.is_right_child():
                node.parent.balance_factor -= 1
            
            if node.parent.balance_factor != 0:
                self.update_balance(node.parent)

上面的update_balance()方法对当前节点及其父节点进行递归调用，重新平衡整棵树的平衡因子。其中rebalance()进行具体的操作。在树上进行一次或多次旋转可以高效的让AVL树发挥作用的同时不损失性能的平衡操作。旋转操作可以分为左旋和右旋。

左旋

左旋操作分为以下步骤：

  • 将右子节点提升为子树的根节点。
  • 将旧的根节点作为新根节点的左子节点。
  • 如果新根节点已经有一个左子节点了，将其作为新左子节点的右子节点。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class AVLTree(BinarySearchTree):
    """
    这里省略了已有的代码
    """
    
    def rotate_left(self, rot_root):
        new_root = rot_root.right_child
        rot_root.right_child = new_root.left_child
        if new_root.left_child is not None:
            new_root.left_child.parent = rot_root
        new_root.parent = rot_root.parent
        if rot_root.is_root():
            self.root = new_root
        else:
            if rot_root.is_left_child():
                rot_root.parent.left_child = new_root
            else:
                rot_root.parent.right_child = new_root

        new_root.left_child = rot_root
        rot_root.parent = new_root
        # 下面的两步平衡因子的计算推导可见书本p208，这是一个简单的推导过程
        rot_root.balance_factor = rot_root.balance_factor + 1 - min(new_root.balance_factor, 0)
        new_root.balance_factor = new_root.balance_factor + 1 + max(rot_root.balance_factor, 0)

右旋

右旋的操作分为以下步骤：

  • 将左子节点提升为子树的根节点。
  • 将旧根节点作为新根节点的右子节点。
  • 如果新根节点已经有了一个右子节点，将其作为新右子节点的左子节点。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class AVLTree(BinarySearchTree):
    """
    这里省略了已有的代码
    """
    
    def rotate_right(self, rot_root):
        new_root = rot_root.left_child
        rot_root.left_child = new_root.right_child
        if new_root.right_child is not None:
            new_root.right_child.parent = rot_root
        new_root.parent = rot_root.parent
        if rot_root.is_root():
            self.root = new_root
        else:
            if rot_root.is_right_child():
                rot_root.parent.right_child = new_root
            else:
                rot_root.parent.left_child = new_root

        new_root.right_child = rot_root
        rot_root.parent = new_root
        # 下面的两步平衡因子的计算推导可见书本p208，这是一个简单的推导过程
        rot_root.balance_factor = rot_root.balance_factor + 1 - min(new_root.balance_factor, 0)
        new_root.balance_factor = new_root.balance_factor + 1 + max(rot_root.balance_factor, 0)

实现再平衡

如果根节点的平衡因子是-2，那么应该左旋，但是得到另一棵失衡的树，其平衡因子是2，右需要右旋，这样进入了死循环。要解决这种情况，需要遵守以下规则：

  • 如果子树需要左旋，首先检查右子树的平衡因子，如果右子树左倾，就对右子树做一次右旋，再围绕原节点做一次左旋；
  • 如果子树需要右旋，首先检查左子树的平衡因子，如果左子树右倾，就对左子树做一次左旋，再围绕原节点做一次右旋。

rebalance方法实现了上述规则，其代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
class AVLTree(BinarySearchTree):
    """
    这里省略了已有的代码
    """
        
    def rebalance(self, node):
        if node.balance_factor < 0:
            if node.right_child.balance_factor > 0:
                self.rotate_right(node.right_child)
                self.rotate_left(node)
            else:
                self.rotate_left(node)
        elif node.balance_factor > 0:
            if node.left_child.balance_factor < 0:
                self.rotate_left(node.left_child)
                self.rotate_right(node)
            else:
                self.rotate_right(node)

实现delete方法

暂未实现

Tags: python 二叉搜索树 平衡二叉搜索树
← 利用二叉堆实现优先级队列 操作系统-分布式系统 →
赏
使用支付宝打赏
使用微信打赏

若你觉得我的文章对你有帮助，欢迎点击上方按钮对我打赏

扫描二维码，分享此文章