在 Python 中构建森林系列：二叉树遍历

Shun Huang

5.00/5 (1投票)

2021年3月26日

CPOL

13分钟阅读

9118

使用 Python 实现二叉搜索树遍历。

本文是《Python 构建森林系列》的第二篇。在本文中，我们不构建树。相反，我们将实现遍历二叉树的函数。当我们使用数据结构来管理数据时，一个重要的用例是遍历我们管理的所有数据。这就是为什么树遍历至关重要的原因。

引言

遍历是一种系统地检查树的节点以访问每个节点一次的方法。遍历二叉树有多种方法：中序、前序、后序和层序。这些名称来自根节点相对于其子树的位置。树遍历也是一种图遍历。中序、前序和后序遍历是深度优先遍历的类型，而层序遍历是广度优先遍历的类型。本文将通过递归和迭代两种方式介绍这些遍历函数的实现。

项目设置

与《构建二叉搜索树》相同，该实现假定使用 Python 3.9 或更高版本。此外，我们在项目中添加了两个模块：用于遍历函数的 traversal.py 和用于其单元测试的 test_traversal.py。添加这两个文件后，我们的项目结构如下所示：

Forest-python
├── forest
│   ├── __init__.py
│   ├── binary_trees
│   │   ├── __init__.py
│   │   ├── binary_search_tree.py
│   │   └── traversal.py
│   └── tree_exceptions.py
└── tests
    ├── __init__.py
    ├── conftest.py
    ├── test_binary_search_tree.py
    └── test_traversal.py

（完整代码可在 forest-python 获取。）

函数接口

我们在此实现的遍历方法力求做到尽可能通用。理论上，只要节点具有左子节点和右子节点字段，我们就可以将遍历方法应用于任何二叉树，事实也确实如此。然而，对于不同类型的二叉树，它们的节点可能包含额外信息。例如，线索化二叉树的节点包含指示该节点是叶子节点还是线索的信息。在这种情况下，当我们遍历线索化二叉树时，我们需要检查额外信息以确定是否到达叶子节点。换句话说，我们不能仅仅检查节点是左右子节点字段是否为空。

我们定义了支持的树类型，而不是为每种树类型的不同条件添加逻辑，这样客户端（即人类或 Linter）就能知道树类型是否受支持。我们支持的类型定义如下：

from typing import Union
from forest.binary_trees import binary_search_tree
SupportedNode = Union[None, binary_search_tree.Node]
SupportedTree = Union[binary_search_tree.BinarySearchTree]

目前，我们支持节点类型 None 和 Node，并且仅支持树类型 BinarySearchTree；当我们实现其他类型的二叉树时，我们会添加更多类型。

有了类型定义，我们就可以利用类型注解来定义我们的遍历函数（以中序遍历为例）。

from typing import Any, Iterator, Optional

Pairs = Iterator[tuple[Any, Any]]
"""Iterator of Key-Value pairs, yield by traversal functions. For type checking"""

def inorder_traverse(tree: SupportedTree, recursive: bool = True) -> Pairs:
    …

请注意，我们还为返回值定义了一个自定义类型 Pairs。尽管这种方式需要编写更多代码，但正如《Python 之禅》所建议的，“显式优于隐式”。这样，客户端就可以清楚地看到正确的输入参数类型以及函数将返回的类型，并且类型检查工具（如 mypy）可以防止类型不匹配问题。此外，我们将为深度优先遍历实现递归和非递归两种方式，这就是为什么我们有第二个参数 recursive。

为什么不使用 @overload？

Python 是一种动态类型语言，而函数重载在 Python 中意义不大。但我们可以利用类型注解来防止类型不匹配问题。

自 Python 3.5 起，PEP-484 引入了 @overload 装饰器。根据官方文档，“@overload 装饰器允许描述支持不同参数类型组合的函数和方法。”听起来很棒。但是，它仅对类型检查器有益。在运行时，客户端代码仍然可以向函数传递任何参数。文档还说：“一个重载的示例，它提供比使用联合类型或类型变量所能表达的更精确的类型。” 因此，我们使用 Union 定义了 SupportedTree 类型。使用 SupportedTree 也比定义多个 @overload 装饰的定义代码量少。

为什么返回类型 Pairs 是一个迭代器？

我们的想法是将遍历函数实现为生成器。这样做的一个主要好处是，当生成器迭代器返回时，它会从上次中断的地方继续。当树非常大时，这种方法可以节省大量内存。遍历函数的客户端代码可以一次处理一个项目。此外，客户端代码将更简单易读。

中序遍历

中序遍历按以下方法访问二叉树，其中每个节点都可以表示为子树的根。

遍历左子树
访问根节点
遍历右子树

下图演示了中序遍历的思想。黄色部分是根节点 23 的左子树，绿色部分是根节点 23 的右子树。节点旁边的小数字（即 1、2、3）表示遍历顺序。

如果二叉树同时也是二叉搜索树，那么二叉搜索树的属性允许我们通过中序遍历产生升序的键值，如图所示。

可以使用递归或辅助堆栈来完成二叉树遍历。如前一节所述，我们实现了递归和非递归两种方式。

def inorder_traverse(tree: SupportedTree, recursive: bool = True) -> Pairs:
    if recursive:
        return _inorder_traverse(node=tree.root)
    return _inorder_traverse_non_recursive(root=tree.root)

由于我们希望隐藏实现细节，我们将实际实现函数（_inorder_traverse 和 _inorder_traverse_non_recursive）命名为以下划线开头，这意味着它们是内部函数。

递归实现

递归遍历是遍历树的最自然方式，因为中序遍历的定义如此，并且可以实现如下：

def _inorder_traverse(node: SupportedNode) -> Pairs:
    if node:
        yield from _inorder_traverse(node.left)
        yield (node.key, node.data)
        yield from _inorder_traverse(node.right)

请注意，我们使用 yield from 来递归调用 _inorder_traverse。这是因为 _inorder_traverse 是一个生成器；要允许一个生成器将其部分操作委托给另一个生成器，我们需要使用 yield from。

非递归实现

关于中序遍历，我们先访问左子树，然后是根节点，最后是右子树。因此，我们先推入右子树，然后推入根节点，然后再访问左子树。推入顺序是因为堆栈是后进先出的数据结构。访问左子树后，我们弹出堆栈以访问根节点，然后再次弹出以访问右子树。重复这些步骤，直到堆栈为空。

要实现中序遍历函数，我们需要一个堆栈来保存稍后将要访问的子树的根节点。进行非递归遍历时，关键在于：1. 何时将节点推入堆栈以及何时从堆栈中弹出节点，以及 2. 在遍历过程中何时生成（即访问）节点。

当正在访问的节点有右子节点时，我们将右子节点推入堆栈，然后将当前节点也推入堆栈。
当节点从堆栈中弹出时，如果该节点没有右子节点，或者其右子节点与堆栈顶部的节点相同，我们就生成该节点。

我们使用 Python 内置的列表作为堆栈来实现非递归中序遍历，因为内置列表具有后进先出的能力：list.append(x) 将一个项添加到列表末尾，list.pop() 从列表中移除并返回最后一项。

def _inorder_traverse_non_recursive(root: SupportedNode) -> Pairs:
    if root is None:
        raise StopIteration

    stack = []
    if root.right:
        stack.append(root.right)
    stack.append(root)
    current = root.left

    while True:

        if current:
            if current.right:
                stack.append(current.right)
                stack.append(current)
                current = current.left
                continue
            stack.append(current)
            current = current.left

        else:  # current is None

            if len(stack) > 0:
                current = stack.pop()

                if current.right is None:
                    yield (current.key, current.data)
                    current = None
                    continue
                else:  # current.right is not None
                    if len(stack) > 0:
                        if current.right == stack[-1]:
                            yield (current.key, current.data)
                            current = stack.pop() if len(stack) > 0 else None
                            continue
                        else:  # current.right != stack[-1]:
                            # This case means there are more nodes on the right
                            # Keep the current and go back to add them.
                            continue

            else:  # stack is empty
                break

反中序遍历

中序遍历在遍历二叉搜索树时会产生升序的排序输出。如果以反向中序遍历二叉树，则排序结果将是降序的。要进行反向中序遍历，我们首先访问右子树，最后访问左子树。

遍历右子树
访问根节点
遍历左子树

反向中序遍历也可以通过递归和非递归方式实现。并且实现方式与中序遍历对称。

def reverse_inorder_traverse(tree: SupportedTree, recursive: bool = True) -> Pairs:
    if recursive:
        return _reverse_inorder_traverse(node=tree.root)
    return _reverse_inorder_traverse_non_recursive(root=tree.root)

递归反向中序遍历

def _reverse_inorder_traverse(node: SupportedNode) -> Pairs:
    if node:
        yield from _reverse_inorder_traverse(node.right)
        yield (node.key, node.data)
        yield from _reverse_inorder_traverse(node.left)

非递归反向中序遍历

def _reverse_inorder_traverse_non_recursive(root: SupportedNode) -> Pairs:
    if root is None:
        raise StopIteration

    stack = []
    if root.left:
        stack.append(root.left)
    stack.append(root)
    current = root.right

    while True:

        if current:
            if current.left:
                stack.append(current.left)
                stack.append(current)
                current = current.right
                continue
            stack.append(current)
            current = current.right

        else:  # current is None

            if len(stack) > 0:
                current = stack.pop()

                if current.left is None:
                    yield (current.key, current.data)
                    current = None
                    continue
                else:  # current.right is not None
                    if len(stack) > 0:
                        if current.left == stack[-1]:
                            yield (current.key, current.data)
                            current = stack.pop() if len(stack) > 0 else None
                            continue
                        else:  # current.right != stack[-1]:
                            # This case means there are more nodes on the right
                            # Keep the current and go back to add them.
                            continue

            else:  # stack is empty
                break

前序遍历

前序遍历按以下方法访问二叉树：

访问根节点。
遍历左子树。
遍历右子树。

下图展示了前序遍历的思想。黄色部分是根节点 23 的左子树，绿色部分是根节点 23 的右子树。与中序遍历相同，节点旁边的小数字（即 1、2、3）表示遍历顺序。

我们也实现了递归和非递归的前序遍历。

def preorder_traverse(tree: SupportedTree, recursive: bool = True) -> Pairs:
    if recursive:
        return _preorder_traverse(node=tree.root)
    return _preorder_traverse_non_recursive(root=tree.root)

递归前序遍历

递归实现很简单，只需按照遍历顺序进行即可。

def _preorder_traverse(node: SupportedNode) -> Pairs:
    if node:
        yield (node.key, node.data)
        yield from _preorder_traverse(node.left)
        yield from _preorder_traverse(node.right)

非递归前序遍历

关于非递归实现，下图展示了非递归前序遍历。

非递归前序遍历比非递归中序遍历简单。由于我们先访问根节点，因此该过程可以视为以下步骤：

当我们访问一个节点时，我们将它的右子节点推入堆栈（如果存在），然后将它的左子节点推入堆栈（如果存在）。
当一个节点从堆栈中弹出时，我们就生成该节点。

def _preorder_traverse_non_recursive(root: SupportedNode) -> Pairs:
    if root is None:
        raise StopIteration

    stack = [root]

    while len(stack) > 0:
        temp = stack.pop()
        yield (temp.key, temp.data)

        # Because stack is FILO, insert right child before left child.
        if temp.right:
            stack.append(temp.right)

        if temp.left:
            stack.append(temp.left)

后序遍历

后序遍历按以下方法访问二叉树：

遍历左子树。
遍历右子树。
访问根节点。

下图展示了后序遍历的思想。与之前的遍历类似，节点旁边的小数字（即 1、2、3）表示遍历顺序。

我们再次通过递归和非递归方法构建后序遍历。

def postorder_traverse(tree: SupportedTree, recursive: bool = True) -> Pairs:
    if recursive:
        return _postorder_traverse(node=tree.root)
    return _postorder_traverse_non_recursive(root=tree.root)

递归后序遍历

与中序和前序遍历类似，递归实现也很简单。

def _postorder_traverse(node: SupportedNode) -> Pairs:
    if node:
        yield from _postorder_traverse(node.left)
        yield from _postorder_traverse(node.right)
        yield (node.key, node.data)

非递归后序遍历

然而，后序遍历的非递归实现有些复杂。下图以非递归方式演示了后序遍历。

与非递归中序遍历类似，当我们执行非递归后序遍历时，如果一个被访问的节点有右子节点，我们就将其右子节点推入堆栈，然后也将当前节点推入堆栈。
当节点从堆栈中弹出时，如果该节点没有右子节点，或者堆栈变为空，我们就生成该节点。此外，如果从堆栈中弹出的节点有右子节点，并且该子节点与堆栈顶部的节点不同，则意味着我们已经访问了右子树。在这种情况下，我们可以生成该节点。否则，我们将堆栈顶部的节点与当前节点交换，然后遍历到右子树。

def _postorder_traverse_non_recursive(root: SupportedNode) -> Pairs:
    if root is None:
        raise StopIteration

    stack = []
    if root.right:
        stack.append(root.right)
    stack.append(root)
    current = root.left

    while True:

        if current:
            if current.right:
                stack.append(current.right)
                stack.append(current)
                current = current.left
                continue
            else:  # current.right is None
                if current.left:
                    stack.append(current)
                else:
                    yield (current.key, current.data)

                current = current.left

        else:  # current is None
            if len(stack) > 0:
                current = stack.pop()

                if current.right is None:
                    yield (current.key, current.data)
                    current = None
                else:  # current.right is not None
                    if len(stack) > 0:
                        if current.right != stack[-1]:
                            yield (current.key, current.data)
                            current = None
                        else:  # current.right == stack[-1]
                            temp = stack.pop()
                            stack.append(current)
                            current = temp

                    else:  # stack is empty
                        yield (current.key, current.data)
                        break
            else:  # stack is empty
                break

层序遍历

与之前的深度优先遍历不同，层序遍历是广度优先的。在这种情况下，我们在进入下一层之前访问完每一层的所有节点。这个思想如下图所示：

我们不使用堆栈，而是使用队列来实现层序遍历函数。队列是先进先出的数据结构。对于每个节点，先访问节点，然后将其子节点放入队列。从队列中出队一个节点，先访问该节点，然后将该节点的子节点入队。重复此过程，直到队列为空。

我们还使用 Python 内置的列表作为队列来实现层序函数，因为内置列表也具有先进先出的能力：list.append(x) 将一个项添加到列表末尾，list.pop(0) 从列表中移除并返回第一项。

def levelorder_traverse(tree: SupportedTree) -> Pairs:
    queue = [tree.root]

    while len(queue) > 0:
        temp = queue.pop(0)
        if temp:
            yield (temp.key, temp.data)
            if temp.left:
                queue.append(temp.left)

            if temp.right:
                queue.append(temp.right)

测试

与往常一样，我们应该尽可能地为我们的代码编写单元测试。在这里，我们使用了我们在《构建二叉搜索树》中创建的 conftest.py 中的 basic_tree 函数来测试我们的遍历函数。在 test_traversal.py 中，我们可以进行如下单元测试：我们检查遍历输出是否与预期相同（以 postorder_traverse 为例）。

def test_binary_search_tree_traversal(basic_tree):
    """Test binary search tree traversal."""
    tree = binary_search_tree.BinarySearchTree()

    for key, data in basic_tree:
        tree.insert(key=key, data=data)

    assert [item for item in traversal.postorder_traverse(tree)] == [
        (1, "1"), (7, "7"), (15, "15"), (22, "22"), (20, "20"), (11, "11"),
        (4, "4"), (24, "24"), (34, "34"), (30, "30"), (23, "23")
    ]

（完整的单元测试请参见 test_traversal.py。）

分析

深度优先遍历：中序、反向中序、前序和后序

中序、反向中序、前序和后序遍历的运行时间

在图算法中，深度优先搜索的运行时间为 O(V+E)，其中 V 是顶点的数量，E 是边的数量（关于运行时间为 O(V+E) 的原因，请参阅《算法导论》第 22.3 章）。由于中序、反向中序、前序和后序遍历是深度优先遍历的类型，我们可以将这些遍历映射到深度优先搜索。因此，我们的运行时间为 O(V+E)。二叉树中的顶点是节点。此外，对于一棵二叉树，如果它有 n 个节点，那么它必须有 n-1 条边。因此，我们可以将运行时间重写为 O(n+(n-1))=O(2n-1)=O(n)，其中 n 是节点数。因此，它们的运行时间变为 O(n)。

关于空间复杂度，我们需要根据我们的方法进行分析：递归或非递归。

递归方法中的空间复杂度

空间使用是当我们采用递归方法时的调用堆栈，即函数调用的深度。我们从根节点到叶子节点遍历二叉树，因此树的高度决定了函数调用的深度。因此，如果一棵二叉树有 n 个节点，我们知道平均情况是 O(lg n)，当树是随机构建的时；最坏情况是 O(n)，当树是线性链式时。

非递归方法中的空间复杂度

如果我们采用非递归方法，空间使用是堆栈的大小。与递归方法相同，堆栈大小由树的高度决定。因此，对于一棵有 n 个节点的二叉树，我们知道平均情况是 O(lg n)，当树是随机构建的时；最坏情况是 O(n)，当树是线性链式时。

广度优先遍历：层序

广度优先搜索算法的时间复杂度为 O(V+E)，其中 V 是顶点的数量，E 是边的数量（有关详细信息，请参阅《算法导论》第 22.2 章）。因此，层序遍历的运行时间也为 O(n)，其中 n 是节点数。原因与上面讨论的中序和其他深度优先类型的遍历相同。

由于我们使用队列来实现层序遍历，因此空间成本是队列的大小。我们在层序遍历中使用队列是为了跟踪同一级别的节点。队列中可能存在节点的最大数量是叶子级别，即 2^h，其中 h 是树的高度（这可以通过数学归纳法轻松证明），队列大小的最坏情况发生在树平衡时，即叶子节点的最大数量。我们也知道平衡树也是树高的最佳情况，即 (lg n)，其中 n 是节点数。因此，队列大小的最坏情况变为 O(2^h)=O(2^(lg n))=O(n)，其中 n 是节点数。

然而，当每个级别只有一个节点时，队列大小的最佳情况发生，即树是线性链式的。在这种情况下，空间复杂度变为常数，即 O(1)。

下表总结了每种遍历实现的复杂性。

示例

既然我们已经实现了遍历例程，我们就可以使用这些遍历函数来实现《二叉搜索树》示例中的 Map 类的 __iter__ 方法，并迭代 map 对象的 data。

from typing import Any, Optional

from forest.binary_trees import binary_search_tree
from forest.binary_trees import traversal

class Map:
    """Key-value Map implemented by a Binary Search Tree."""

    def __init__(self) -> None:
        self._bst = binary_search_tree.BinarySearchTree()

    def __setitem__(self, key: Any, value: Any) -> None:
        """Insert (key, value) item into the map."""
        self._bst.insert(key=key, data=value)

    def __getitem__(self, key: Any) -> Optional[Any]:
        """Get the data by the given key."""
        node = self._bst.search(key=key)
        if node:
            return node.data
        return None

    def __delitem__(self, key: Any) -> None:
        """Remove a (key, value) pair from the map."""
        self._bst.delete(key=key)

    def __iter__(self) -> traversal.Pairs:
        """Iterate the data in the map."""
        return traversal.inorder_traverse(tree=self._bst)

    @property
    def empty(self) -> bool:
        """Return `True` if the map is empty; `False` otherwise."""
        return self._bst.empty

if __name__ == "__main__":

    # Initialize the Map instance.
    contacts = Map()

    # Add some items.
    contacts["Mark"] = "mark@email.com"
    contacts["John"] = "john@email.com"
    contacts["Luke"] = "luke@email.com"

    # Retrieve an email
    print(contacts["Mark"])

    # Delete one item.
    del contacts["John"]

    # Check the deleted item.
    print(contacts["John"])  # This will print None

    # Iterate the items.
    for contact in contacts:
        print(contact)

（完整示例可在 bst_map.py 中找到。）

摘要

二叉树遍历是系统访问二叉树中每个节点的方法，每个节点仅访问一次。它们是图搜索算法（例如深度优先搜索和广度优先搜索）的特例，可以通过不同的方法实现，每种方法都有其优缺点。接下来的文章将介绍改进的二叉搜索树——线索化二叉搜索树，它允许我们在不使用堆栈或递归方法的情况下遍历树，即遍历的空间复杂度始终是常数。

（最初发布于 Shun's Vineyard，发布日期为 2021 年 3 月 17 日。）