完成 BST 操作和时间复杂度

Author: PegasusWang

docs/17_二叉查找树/binary_search_tree.md

@@ -2,7 +2,7 @@
 
 二叉树的一种应用就是来实现堆，今天我们再看看用二叉查找树。
 前面有章节说到了查找操作，包括线性查找和二分查找，线性查找效率比较低，二分又要求必须是有序的序列，
-为了维持有序插入的代价比较高。能不能有一种插入和查找都比较快的数据结构呢？
+为了维持有序插入的代价比较高。能不能有一种插入和查找都比较快的数据结构呢？二叉查找树就是这样一种结构，可以高效地插入和查询节点。
 
 # BST 定义
 
@@ -117,10 +117,19 @@ bst = BST.build_from(NODE_LIST)
 
 ## 插入
 插入节点的时候我们需要一直保持 BST 的性质，每次插入一个节点，我们都通过递归比较把它放到正确的位置。
+你会发现新节点总是被作为叶子结点插入。（请你思考这是为什么）
+
+![](./bst_insert.png)
 
 ```py
     def _bst_insert(self, subtree, key, value):
-        if subtree is None:
+        """ 插入并且返回根节点
+
+        :param subtree:
+        :param key:
+        :param value:
+        """
+        if subtree is None:   # 插入的节点一定是根节点，包括 root 为空的情况
             subtree = BSTNode(key, value)
         elif key < subtree.key:
             subtree.left = self._bst_insert(subtree.left, key, value)
@@ -139,11 +148,95 @@ bst = BST.build_from(NODE_LIST)
             return True
 ```
 
-## 删除
+## 删除节点
+删除操作相比上边的操作要麻烦很多，首先需要定位一个节点，删除节点后，我们需要始终保持 BST 的性质。
+删除一个节点涉及到三种情况：
+
+- 节点是叶节点
+- 节点有一个孩子
+- 节点有两个孩子
+
+我们分别来看看三种情况下如何删除一个节点：
+
+#### 删除叶节点
+这是最简单的一种情况，只需要把它的父亲指向它的指针设置为 None 就好。
+
+![](./bst_remove_leaf.png)
+
+#### 删除只有一个孩子的节点
+删除有一个孩子的节点时，我们拿掉需要删除的节点，之后把它的父亲指向它的孩子就行，以为根据 BST
+左子树都小于节点，右子树都大于节点的特性，删除它之后这个条件依旧满足。
+
+![](./bst_remove_node_with_one_child.png)
+
+#### 删除有两个孩子的内部节点
+假如我们想删除 12 这个节点改怎么做呢？你的第一反应可能是按照下图的方式：
+
+![](./remove_interior_replace.png)
+
+但是这种方式可能会影响树的高度，降低查找的效率。这里我们用另一种非常巧妙的方式。
+还记得上边提到的吗，如果你中序遍历 BST 并且输出每个节点的 key，你会发现就是一个有序的数组。
+[1 4 12 23 29 37 41 60 71 84 90 100]。 这里我们定义两个概念，逻辑前任(predecessor)和后继(successor)，请看下图:
+
+![](./predecessor_successor.png)
+
+12 在中序遍历中的逻辑前任和后继分别是 4 和 23 节点。于是我们还有一种方法来删除 12 这个节点：
+
+- 找到节点待删除节点 N(12) 的后继节点 S(23)
+- 复制节点 S 到节点 N
+- 删除节点 S
+
+说白了就是找到后继并且替换，这里之所以能保证这种方法是正确的，你会发现替换后依旧是保持了 BST 的性质。
+有个问题是如何找到后继节点呢？待删除节点的右子树的最小的节点不就是后继嘛，上边我们已经实现了找到最小 key 的方法了。
+
+
+![](./find_successor.png)
+
+我们开始编写代码实现，和之前的操作类似，我们还是通过辅助函数的形式来实现，这个递归函数会比较复杂，请你仔细理解:
+
+```py
+    def _bst_remove(self, subtree, key):
+        """删除节点并返回根节点"""
+        if subtree is None:
+            return None
+        elif key < subtree.key:
+            subtree.left = self._bst_remove(subtree.left, key)
+            return subtree
+        elif key > subtree.key:
+            subtree.right = self._bst_remove(subtree.right, key)
+            return subtree
+        else:  # 找到了需要删除的节点
+            if subtree.left is None and subtree.right is None:    # left node
+                return None
+            elif subtree.left is None or subtree.right is None:  # 只有一个孩子
+                if subtree.left is not None:
+                    return subtree.left
+                else:
+                    return subtree.right
+            else:  # 俩孩子
+                successor_node = self._bst_min_node(subtree.right)
+                subtree.key, subtree.value = successor_node.key, subtree.value
+                subtree.right = self._bst_remove(subtree.right, successor_node.key)
+                return subtree
+
+    def remove(self, key):
+        assert key in self
+        self.size -= 1
+```
+
+完整代码你可以在本章的 bst.py  找到。
+
+# 时间复杂度分析
+
+上边介绍的操作时间复杂度和二叉树的形状有关。平均来说时间复杂度是和树的高度成正比的，树的高度 h 是 log(n)，
+但是最坏情况下以上操作的时间复杂度都是 O(n)。为了改善 BST 有很多变种，感兴趣请参考延伸阅读中的内容。
+
+![](./bst_worstcase.png)
+
 
 # 练习题：
 - 请你实现查找 BST 最大值的函数
 
 
 # 延伸阅读
-- 《Data Structures and Algorithms in Python》14 章
+- 《Data Structures and Algorithms in Python》14 章，树的概念和算法还有很多，我们这里介绍最基本的

docs/17_二叉查找树/bst.py

@@ -38,6 +38,10 @@ def _bst_search(self, subtree, key):
         else:
             return subtree
 
+    def __contains__(self, key):
+        """实现 in 操作符"""
+        return self._bst_search(self.root, key) is not None
+
     def get(self, key, default=None):
         node = self._bst_search(self.root, key)
         if node is None:
@@ -58,7 +62,13 @@ def bst_min(self):
         return node.value if node else None
 
     def _bst_insert(self, subtree, key, value):
-        if subtree is None:
+        """ 插入并且返回根节点
+
+        :param subtree:
+        :param key:
+        :param value:
+        """
+        if subtree is None:   # 插入的节点一定是根节点，包括 root 为空的情况
             subtree = BSTNode(key, value)
         elif key < subtree.key:
             subtree.left = self._bst_insert(subtree.left, key, value)
@@ -76,6 +86,35 @@ def add(self, key, value):
             self.size += 1
             return True
 
+    def _bst_remove(self, subtree, key):
+        """删除节点并返回根节点"""
+        if subtree is None:
+            return None
+        elif key < subtree.key:
+            subtree.left = self._bst_remove(subtree.left, key)
+            return subtree
+        elif key > subtree.key:
+            subtree.right = self._bst_remove(subtree.right, key)
+            return subtree
+        else:  # 找到了需要删除的节点
+            if subtree.left is None and subtree.right is None:    # left node
+                return None
+            elif subtree.left is None or subtree.right is None:  # 只有一个孩子
+                if subtree.left is not None:
+                    return subtree.left
+                else:
+                    return subtree.right
+            else:  # 俩孩子
+                successor_node = self._bst_min_node(subtree.right)
+                subtree.key, subtree.value = successor_node.key, subtree.value
+                subtree.right = self._bst_remove(subtree.right, successor_node.key)
+                return subtree
+
+    def remove(self, key):
+        assert key in self
+        self.size -= 1
+        return self._bst_remove(self.root, key)
+
 
 NODE_LIST = [
     {'key': 60, 'left': 12, 'right': 90, 'is_root': True},
@@ -105,3 +144,12 @@ def test_bst_tree():
 
     bst.add(0, 0)
     assert bst.bst_min() == 0
+
+    bst.remove(12)
+    assert bst.get(12) is None
+
+    bst.remove(1)
+    assert bst.get(1) is None
+
+    bst.remove(29)
+    assert bst.get(29) is None