二叉搜索树

我们已经学习了两种从集合中获取键-值对的方法。回想一下,我们讨论过映射抽象数据类型的两种实现,它们分别是列表二分搜索和散列表。本节将探讨二叉搜索树,它是映射的另一种实现。我们感兴趣的不是元素在树中的确切位置,而是如何利用二叉树结构提供高效的搜索。

搜索树的操作

在实现搜索树之前,我们来复习一下映射抽象数据类型提供的接口。你会发现,这个接口类似于 Python 字典。

  • Map() 新建一个空的映射。

  • put(key, val) 往映射中加入一个新的键-值对。如果键已经存在,就用新值替换旧值。

  • get(key) 返回 key 对应的值。如果 key 不存在,则返回 None。

  • del 通过 del map[key] 这样的语句从映射中删除键-值对。

  • len() 返回映射中存储的键-值对的数目。

  • in 通过 key in map 这样的语句,在键存在时返回 True,否则返回 False。

搜索树的实现

二叉搜索树依赖于这样一个性质:小于父节点的键都在左子树中,大于父节点的键则都在右子树中。我们称这个性质为二叉搜索性,它会引导我们实现上述映射接口。图 6-20 描绘了二叉搜索树的这个性质,图中只展示了键,没有展示对应的值。注意,每一对父节点和子节点都具有这个性质。左子树的所有键都小于根节点的键,右子树的所有键则都大于根节点的键。

image 2023 12 13 19 01 06 241
Figure 1. 图6-20 简单的二叉搜索树

接下来看看如何构造二叉搜索树。图 6-20 中的节点是按如下顺序插入键之后形成的:70、31、93、94、14、23、73。因为 70 是第一个插入的键,所以是根节点。31 小于 70,所以成为 70 的左子节点。93 大于 70,所以成为 70 的右子节点。现在树的两层已经满了,所以下一个键会成为 31 或 93 的子节点。94 比 70 和 93 都要大,所以它成了 93 的右子节点。同理,14 比 70 和 31 都要小,所以它成了 31 的左子节点。23 也小于 31,所以它必定在 31 的左子树中。而它又大于 14,所以成了 14 的右子节点。

我们将采用 “节点与引用” 表示法实现二叉搜索树,它类似于我们在实现链表和表达式树时采用的方法。不过,由于必须创建并处理一棵空的二叉搜索树,因此我们将使用两个类。一个称作 BinarySearchTree,另一个称作 TreeNode。BinarySearchTree 类有一个引用,指向作为二叉搜索树根节点的 TreeNode 类。大多数情况下,外面这个类的方法只是检查树是否为空。如果树中有节点,请求就被发往 BinarySearchTree 类的私有方法,这个方法以根节点作为参数。当树为空,或者想删除根节点的键时,需要采取特殊措施。代码清单 6-23 是 BinarySearchTree 类的构造方法及一些其他的方法。

代码清单6-23 BinarySearchTree类
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__()

TreeNode 类提供了很多辅助函数,这大大地简化了 BinarySearchTree 类的工作。代码清单 6-24 是 TreeNode 类的构造方法以及辅助函数。可以看到,很多辅助函数有助于根据子节点的位置(是左还是右)以及自己的子节点类型来给节点归类。

代码清单6-24 TreeNode类
class TreeNode:
   def __init__(self,key,val,left=None,right=None,
                                       parent=None):
        self.key = key
        self.payload = val
        self.leftChild = left
        self.rightChild = right
        self.parent = parent

    def hasLeftChild(self):
        return self.leftChild

    def hasRightChild(self):
        return self.rightChild

    def isLeftChild(self):
        return self.parent and self.parent.leftChild == self

    def isRightChild(self):
        return self.parent and self.parent.rightChild == self

    def isRoot(self):
        return not self.parent

    def isLeaf(self):
        return not (self.rightChild or self.leftChild)

    def hasAnyChildren(self):
        return self.rightChild or self.leftChild

    def hasBothChildren(self):
        return self.rightChild and self.leftChild

    def replaceNodeData(self,key,value,lc,rc):
        self.key = key
        self.payload = value
        self.leftChild = lc
        self.rightChild = rc
        if self.hasLeftChild():
            self.leftChild.parent = self
        if self.hasRightChild():
            self.rightChild.parent = self

TreeNode 类与 6.4.2 节中的 BinaryTree 类有一个很大的区别,那就是显式地将每个节点的父节点记录为它的一个属性。在讨论 del 操作的实现时,你会看到这一点为何重要。

在 TreeNode 类的实现中,另一个有趣之处是使用 Python 的可选参数。可选参数使得在多种环境下创建 TreeNode 更方便。有时,我们想构造一个已有 parent 和 child 的 TreeNode。可以将父节点和子节点作为参数传入。其他时候,只通过键-值对创建 TreeNode,而不传入 parent 和 child。在这种情况下,可选参数使用默认值。

现在有了 BinarySearchTree 和 TreeNode,是时候写一个帮我们构建二叉搜索树的 put 方法了。put 是 BinarySearchTree 类的一个方法。它检查树是否已经有根节点,若没有,就创建一个 TreeNode,并将其作为树的根节点;若有,就调用私有的递归辅助函数 _put,并根据以下算法在树中搜索。

  • 从根节点开始搜索二叉树,比较新键与当前节点的键。如果新键更小,搜索左子树。如果新键更大,搜索右子树。

  • 当没有可供搜索的左(右)子节点时,就说明找到了新键的插入位置。

  • 向树中插入一个节点,做法是创建一个 TreeNode 对象,并将其插入到前一步发现的位置上。

向树中插入新节点的方法如代码清单 6-25 所示。按照上述步骤,我们将 _put 写成递归函数。注意,在向树中插入新的子节点时,currentNode 被作为父节点传入新树。

代码清单6-25 为二叉搜索树插入新节点
def put(self,key,val):
    if self.root:
        self._put(key,val,self.root)
    else:
        self.root = TreeNode(key,val)
    self.size = self.size + 1

def _put(self,key,val,currentNode):
    if key < currentNode.key:
        if currentNode.hasLeftChild():
               self._put(key,val,currentNode.leftChild)
        else:
               currentNode.leftChild = TreeNode(key,val,parent=currentNode)
    else:
        if currentNode.hasRightChild():
               self._put(key,val,currentNode.rightChild)
        else:
               currentNode.rightChild = TreeNode(key,val,parent=currentNode)

插入方法有个重要的问题:不能正确地处理重复的键。遇到重复的键时,它会在已有节点的右子树中创建一个具有同样键的节点。这样做的结果就是搜索时永远发现不了较新的键。要处理重复键插入,更好的做法是用关联的新值替换旧值。这个修复工作留作练习。

定义 put 方法后,就可以方便地通过让 __setitem__ 方法调用 put 方法来重载 [] 运算符。如此一来,就可以写出像 myZipTree['Plymouth'] = 55446 这样的 Python 语句,就如同访问 Python 字典一样。__setitem__ 方法如代码清单 6-26 所示。

代码清单6-26 __setitem__ 方法
def __setitem__(self,k,v):
    self.put(k,v)

图6-21展示了向二叉搜索树中插入新节点的过程。浅灰色节点表示在插入过程中被访问过的节点。

image 2023 12 13 19 09 05 550
Figure 2. 图6-21 插入键为19的新节点

构造出树后,下一个任务就是实现为给定的键取值。get 方法比 put 方法还要简单,因为它只是递归地搜索二叉树,直到访问到叶子节点或者找到匹配的键。在后一种情况下,它会返回节点中存储的值。

get、_get__getitem__ 的实现如代码清单 6-27 所示。 _get 方法中的搜索代码和 _put 方法中选择左右子节点的逻辑相同。注意, _get 方法返回一个 TreeNode 给 get。这样一来,对于其他 BinarySearchTree 方法来说,如果需要使用 TreeNode 有效载荷之外的数据,_get 可以作为灵活的辅助函数使用。

通过实现 __getitem__ 方法,可以写出类似于访问字典的 Python 语句——而实际上使用的是二叉搜索树——比如 z = myZipTree['Fargo']。从代码清单 6-27 可以看出,__getitem__ 方法要做的就是调用 get 方法。

代码清单 6-27查找键对应的值
def get(self,key):
    if self.root:
        res = self._get(key,self.root)
        if res:
               return res.payload
        else:
               return None
    else:
        return None

def _get(self,key,currentNode):
    if not currentNode:
        return None
    elif currentNode.key == key:
        return currentNode
    elif key < currentNode.key:
        return self._get(key,currentNode.leftChild)
    else:
        return self._get(key,currentNode.rightChild)

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

利用 get 方法,可以通过为 BinarySearchTree 编写 __contains__ 方法来实现 in 操作。__contains__ 方法只需调用 get 方法,并在 get 方法返回一个值时返回 True,或在 get 方法返回 None 时返回 False。代码清单6-28 实现了 __contains__ 方法。

代码清单6-28 检查树中是否有某个键
def __contains__(self,key):
    if self._get(key,self.root):
        return True
    else:
        return False

你应该记得,__contains__ 方法重载了 in 运算符,因此我们可以写出这样的语句:

if 'Northfield' in myZipTree:
    print("oom ya ya")

最后,我们将注意力转向二叉搜索树中最有挑战性的方法——删除一个键。第一个任务是在树中搜索并找到要删除的节点。如果树中不止一个节点,使用 _get 方法搜索,找到要移除的 TreeNode。如果树中只有一个节点,则意味着要移除的是根节点,不过仍要确保根节点的键就是要删除的键。无论哪种情况,如果找不到要删除的键,delete 方法都会抛出一个异常,如代码清单6-29 所示。

代码清单6-29 delete方法
def delete(self,key):
   if self.size > 1:
      nodeToRemove = self._get(key,self.root)
      if nodeToRemove:
          self.remove(nodeToRemove)
          self.size = 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 = self.size - 1
   else:
      raise KeyError('Error, key not in tree')

def __delitem__(self,key):
    self.delete(key)

一旦找到待删除键对应的节点,就必须考虑 3 种情况。

(1) 待删除节点没有子节点(如图 6-22 所示)。

image 2023 12 13 19 20 43 985
Figure 3. 图6-22 待删除节点16没有子节点

(2) 待删除节点只有一个子节点(如图 6-23 所示)。

image 2023 12 13 19 21 13 856
Figure 4. 图6-23 待删除节点25有一个子节点

(3) 待删除节点有两个子节点(如图 6-24 所示)。

image 2023 12 13 19 21 42 666
Figure 5. 图6-24 待删除节点5有两个子节点

情况 1 很简单。如果当前节点没有子节点,要做的就是删除这个节点,并移除父节点对这个节点的引用,如代码清单 6-30 所示。

代码清单6-30 情况1:待删除节点没有子节点
if currentNode.isLeaf():
    if currentNode == currentNode.parent.leftChild:
        currentNode.parent.leftChild = None
    else:
        currentNode.parent.rightChild = None

情况2 稍微复杂些。如果待删除节点只有一个子节点,那么可以用子节点取代待删除节点,如代码清单6-31 所示。查看这段代码后会发现,它考虑了 6 种情况。由于左右子节点的情况是对称的,因此只需要讨论当前节点有左子节点的情况。

  • 如果当前节点是一个左子节点,只需将当前节点的左子节点对父节点的引用改为指向当前节点的父节点,然后将父节点对当前节点的引用改为指向当前节点的左子节点。

  • 如果当前节点是一个右子节点,只需将当前节点的右子节点对父节点的引用改为指向当前节点的父节点,然后将父节点对当前节点的引用改为指向当前节点的右子节点。

  • 如果当前节点没有父节点,那它肯定是根节点。调用 replaceNodeData 方法,替换根节点的 key、payload、leftChild 和 rightChild 数据。

代码清单6-31 情况2:待删除节点只有一个子节点
else: # this node has one child
   if currentNode.hasLeftChild():
      if currentNode.isLeftChild():
          currentNode.leftChild.parent = currentNode.parent
          currentNode.parent.leftChild = currentNode.leftChild
      elif currentNode.isRightChild():
          currentNode.leftChild.parent = currentNode.parent
          currentNode.parent.rightChild = currentNode.leftChild
      else:
          currentNode.replaceNodeData(currentNode.leftChild.key,
                             currentNode.leftChild.payload,
                             currentNode.leftChild.leftChild,
                             currentNode.leftChild.rightChild)
   else:
      if currentNode.isLeftChild():
          currentNode.rightChild.parent = currentNode.parent
          currentNode.parent.leftChild = currentNode.rightChild
      elif currentNode.isRightChild():
          currentNode.rightChild.parent = currentNode.parent
          currentNode.parent.rightChild = currentNode.rightChild
      else:
          currentNode.replaceNodeData(currentNode.rightChild.key,
                             currentNode.rightChild.payload,
                             currentNode.rightChild.leftChild,
                             currentNode.rightChild.rightChild)

情况3 最难处理。如果一个节点有两个子节点,那就不太可能仅靠用其中一个子节点取代它来解决问题。不过,可以搜索整棵树,找到可以替换待删除节点的节点。候选节点要能为左右子树都保持二叉搜索树的关系,也就是树中具有次大键的节点。我们将这个节点称为后继节点,有一种方法能快速找到它。后继节点的子节点必定不会多于一个,所以我们知道如何按照已实现的两种删除方法来移除它。移除后继节点后,只需直接将它放到树中待删除节点的位置上即可。

处理情况3的代码如代码清单 6-32 所示。注意,我们用辅助函数 findSuccessor 和 findMin 来寻找后继节点,并用 spliceOut 方法移除它(如代码清单 6-34 所示)。之所以用 spliceOut 方法,是因为它可以直接访问待拼接的节点,并进行正确的修改。虽然也可以递归调用 delete,但那样做会浪费时间重复搜索键的节点。

代码清单6-32 情况3:待删除节点有两个子节点
elif currentNode.hasBothChildren(): #interior
        succ = currentNode.findSuccessor()
        succ.spliceOut()
        currentNode.key = succ.key
        currentNode.payload = succ.payload

寻找后继节点的代码如代码清单 6-33 所示。可以看出,这是 TreeNode 类的一个方法。它利用的二叉搜索树属性,也是从小到大打印出树节点的中序遍历所利用的。在查找后继节点时,要考虑以下 3 种情况。

  1. 如果节点有右子节点,那么后继节点就是右子树中最小的节点。

  2. 如果节点没有右子节点,并且其本身是父节点的左子节点,那么后继节点就是父节点。

  3. 如果节点是父节点的右子节点,并且其本身没有右子节点,那么后继节点就是除其本身外父节点的后继节点。

在试图从一棵二叉搜索树中删除节点时,上述第一个条件是唯一重要的。但是,findSuccessor 方法还有其他用途,本章末会进行探索。

findMin 方法用来查找子树中最小的键。可以确定,在任意二叉搜索树中,最小的键就是最左边的子节点。鉴于此,findMin 方法只需沿着子树中每个节点的leftChild 引用走,直到遇到一个没有左子节点的节点。代码清单 6-35 给出了完整的 remove 方法。

代码清单6-33 寻找后继节点
def findSuccessor(self):
    succ = None
    if self.hasRightChild():
        succ = self.rightChild.findMin()
    else:
        if self.parent:
               if self.isLeftChild():
                   succ = self.parent
               else:
                   self.parent.rightChild = None
                   succ = self.parent.findSuccessor()
                   self.parent.rightChild = self
    return succ

def findMin(self):
    current = self
    while current.hasLeftChild():
        current = current.leftChild
    return current
代码清单6-34 spliceOut方法
def spliceOut(self):
    if self.isLeaf():
        if self.isLeftChild():
               self.parent.leftChild = None
        else:
               self.parent.rightChild = None
    elif self.hasAnyChildren():
        if self.hasLeftChild():
               if self.isLeftChild():
                  self.parent.leftChild = self.leftChild
               else:
                  self.parent.rightChild = self.leftChild
               self.leftChild.parent = self.parent
        else:
               if self.isLeftChild():
                  self.parent.leftChild = self.rightChild
               else:
                  self.parent.rightChild = self.rightChild
               self.rightChild.parent = self.parent
image 2023 12 13 19 29 17 500
Figure 6. 代码清单6-35 remove方法

现在来看看最后一个二叉搜索树接口方法。假设我们想按顺序遍历树中的键。我们在字典中就是这么做的,为什么不在树中试试呢?我们已经知道如何按顺序遍历二叉树——使用中序遍历算法。不过,为了创建迭代器,还需要做更多工作,因为迭代器每次调用只返回一个节点。

Python 为创建迭代器提供了一个很强大的函数,即 yield。与 return 类似,yield 每次向调用方返回一个值。除此之外,yield 还会冻结函数的状态,因此下次调用函数时,会从这次离开之处继续。创建可迭代对象的函数被称作 生成器

二叉搜索树迭代器的代码如代码清单6-36 所示。请仔细看看这份代码。乍看之下,你可能会认为它不是递归的。但是,因为 __iter__ 重载了循环的 for x in 操作,所以它真的是递归的!由于在 TreeNode 实例上递归,因此 __iter__ 方法被定义在 TreeNode 类中。

代码清单6-36 二叉搜索树迭代器
def __iter__(self):
   if self:
      if self.hasLeftChild():
             for elem in self.leftChiLd:
                yield elem
      yield self.key
      if self.hasRightChild():
             for elem in self.rightChild:
                yield elem

搜索树的分析

至此,我们已经完整地实现了二叉搜索树,接下来简单地分析它的各个方法。先分析 put 方法,限制其性能的因素是二叉树的高度。6.3 节曾说过,树的高度是其中节点层数的最大值。高度之所以是限制因素,是因为在搜索合适的插入位置时,每一层最多需要做一次比较。

那么,二叉树的高度是多少呢?答案取决于键的插入方式。如果键的插入顺序是随机的,那么树的高度约为 \$log_2n\$,其中 n 为树的节点数。这是因为,若键是随机分布的,那么小于和大于根节点的键大约各占一半。二叉树的顶层有 1 个根节点,第 1 层有 2 个节点,第 2 层有 4 个节点,依此类推。在完全平衡的二叉树中,节点总数是 \$2^(h+1)-1\$,其中 h 代表树的高度。

在完全平衡的二叉树中,左右子树的节点数相同。最坏情况下,put 的时间复杂度是 \$O(log_2n)\$,其中 n 是树的节点数。注意,这是上一段所述运算的逆运算。所以,\$log_2n\$ 是树的高度,代表 put 在搜索合适的插入位置时所需的最大比较次数。

不幸的是,按顺序插入键可以构造出一棵高度为 n 的搜索树!图 6-25 就是一个例子,这时 put 方法的时间复杂度为 \$O(n)\$。

image 2023 12 13 19 35 21 393
Figure 7. 图6-25 偏斜的二叉搜索树

既然理解了为何说 put 的性能由树的高度决定,你应该可以猜到,get、in 和 del 也都如此。get 在树中查找键,最坏情况就是沿着树一直搜到底也没找到。乍看之下,del 可能更复杂,因为在删除节点前可能还得找到后继节点。但是查找后继节点的最坏情况也受限于树的高度,也就是把工作量加一倍。所以,对于不平衡的树来说,最坏情况下的时间复杂度仍是 \$O(n)\$。