二叉搜索树的python实现（递归调用生成器完成迭代，实现for循环升序打印）

目录

• 二叉搜索树的基本性质
• 二叉搜索树增，删，查的实现
• 递归调用生成器实现for循环的升序迭代
• 补充：可迭代对象，迭代器，生成器

二叉搜索树的基本性质

• 每个节点的左节点的值小于该节点，右节点的值大于该节点

• 操作	时间复杂度
  增加	O(logn)
  删除	O(logn)
  查找	O(logn)

  上述复杂度只是该二叉树是完全二叉树的理想情况，实际上会随着树结构变得稀疏而变差，可以用平衡二叉树来改善

二叉搜索树增删查的实现

该数据结构，作为二分搜索，哈希表外的第三种高效的搜索方式，存储的数据是key-value类型的 

，在实现这个数据结构时，使用“节点与引用“的方式，和链表的实现类似，首先实现基本的节点类

为了方便后续代码阅写了一些辅助函数

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

    def has_left_child(self):
        return self.left_child

    def has_right_child(self):
        return self.right_child

    def has_parent(self):
        return self.parent is not None

    def is_left_child(self):
        return self.has_parent() and self.parent.left_child == self

    def is_right_child(self):
        return self.has_parent() and self.parent.right_child == self

    def is_root(self):
        return not self.has_parent()

    def is_leaf(self):
        return self.left_child is None and self.right_child is None

    def has_any_children(self):
        return self.left_child is not None or self.right_child is not None

    def has_both_children(self):
        return self.left_child is not None and self.right_child is not None

接着要实现二叉搜索树的类

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

    def length(self):
        return self.size

    def __len__(self):
        return self.size

添加

从根节点开始，比当前节点小，则搜查左子树，比当前节点大，则继续搜索右子树，直到当前节点的左子节点或右子节点为None，就将增加的值添加到这里。如果搜索过程中出现相等的情况，那么将新添加的数据替换树中原有的数据

重写了__setitem__方法是为了能够以x[a] = b 的方式添加数据

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

    def _put(self, key, val, cur_node):
        if key == cur_node.key:
            cur_node.key = key
            cur_node.value = val
        elif key < cur_node.key:
            if cur_node.has_left_child():
                self._put(key, val, cur_node.left_child)
            else:
                node = TreeNode(key, val, parent=cur_node)
                cur_node.left_child = node
        else:
            if cur_node.has_right_child():
                self._put(key, val, cur_node.right_child)
            else:
                node = TreeNode(key, val, parent=cur_node)
                cur_node.right_child = node

    def __setitem__(self, key, value):
        self.put(key, value)

查找

查找和添加相同，都是自顶向下。如果有数据的key和查找的key相同，就返回该数据的value，没有的话就返回None

重写__getitem__方法是为了能够以x[a]的方式访问数据的value

   def get(self, key):
        if self.size == 0:
            return None
        else:
            res = self._get(key, self.root)
            if res is not None:
                return res.value
            else:
                return None

    def _get(self, key, cur_node):
        if key == cur_node.key:
            return cur_node
        elif key < cur_node.key:
            if cur_node.has_left_child():
                return self._get(key, cur_node.left_child)
            else:
                return None
        else:
            if cur_node.has_right_child():
                return self._get(key, cur_node.right_child)
            else:
                return None

    def __getitem__(self, item):
        return self.get(item)

删除

删除共有三种情况

1. 待删除节点是叶节点：直接将该节点的父节点的左/右子节点改为None
2. 待删除节点有一个节点：共分为六种情况，具体看代码（很简单）
3. 待删除节点有两个节点：为了使删除该节点后树仍然满足搜索树的特征，需要后继节点，该后继节点需要比待删除节点的左子树上的节点的value都大，比右子树上的节点的value都小，那么需要找的就是 a.待删除节点左子树上最右边的节点 b.待删除节点右子树上最左边的节点

下面代码寻找后继节点的方式是寻找待删除节点右子树的最左边的节点

至于为什么要分成两个函数来实现，用户在调用删除方法时传入的肯定是待删除的key而不是节点吧😂所以只好分成两个部分了

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

    def remove(self, cur_node):
        if cur_node.is_leaf():
            # 叶节点
            if cur_node.is_left_child():
                cur_node.parent.left_child = None
            else:
                cur_node.parent.right_child = None
        else:
            # 只有一个节点 共有六种情况
            if not cur_node.has_both_children():
                if cur_node.is_left_child():
                    if cur_node.has_left_child():
                        cur_node.parent.left_child = cur_node.left_child
                        cur_node.left_child.parent = cur_node.parent
                    else:
                        cur_node.parent.left_child = cur_node.right_child
                        cur_node.right_child.parent = cur_node.parent
                elif cur_node.is_right_child():
                    if cur_node.has_left_child():
                        cur_node.parent.right_child = cur_node.left_child
                        cur_node.left_child.parent = cur_node.parent
                    else:
                        cur_node.parent.right_child = cur_node.right_node
                        cur_node.right_node.parent = cur_node.parent
                else:
                    if cur_node.has_left_child():
                        cur_node.left_child = None
                    else:
                        cur_node.right_child = None
            else:
                # 有左右两个子节点，此处寻找右子树的最左侧的节点
                successor_node = cur_node.right_child
                while successor_node.has_left_child():
                    successor_node = successor_node.left_child
                self.remove(successor_node)
                cur_node.key = successor_node.key
                cur_node.value = successor_node.value

递归调用生成器实现for循环的升序迭代

• 根据搜索树的特点，为了实现升序，需要进行中序遍历
• 中序遍历的过程中又需要对中序遍历函数递归调用
• 为了实现for循环，需要将二叉搜索树构造为可迭代对象，也就是需要重写类的魔术方法__iter__，同时该方法需要返回一个迭代器（可以用yield将__iter__方法本人变成生成器，也相当于一个迭代器）
• 综上，将中序遍历的代码写在__iter__中，并将return换为yield即可
• 需要注意的是：虽然生成器的定义方法和函数一样，但是使用方法和类一样。例如一个生成器函数名字叫generator，使用a = generator()相当于实例化了一个生成器，但还并没有开始执行该生成器。但如果是函数的话，a就是函数的返回值
• 所以在递归中使用return 函数名()的递归方式，在生成器中需要改为使用for循环，因为如果只是yield 生成器()的话，只是新实例化了一个生成器就直接返回了，并没有往下继续执行。

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

    # self._preorder()作为一个生成器（是特殊的迭代器），用于中序遍历二叉搜索树从而实现升序打印。
    # 为什么要把preorder写在iter之外：因为递归调用的过程中需要知道从哪个节点开始的，但是iter函数不允许传参数
    def _preorder(self, cur_node):
        if cur_node.has_left_child():
            for i, j in self._preorder(cur_node.left_child):
                yield i, j
        yield cur_node.key, cur_node.value
        if cur_node.has_right_child():
            for i, j in self._preorder(cur_node.right_child):
                yield i, j