在 Python 中构建森林系列：AVL 树 vs 红黑树

指标

我们将要构建的 Metrics 库受到了 Metrics 项目的启发，它有两个主要组件：MetricRegistry 和受支持的指标类型，包括 Counter 和 Histogram。

MetricRegistry

MetricRegistry 是使用 Metrics 库的起点，它是用于跟踪应用程序的所有指标的集合。因此，例如，如果我们想跟踪一项操作，比如旋转次数，我们可以定义一个 Counter 类型的指标来监视旋转次数，并将 Counter 指标注册到 MetricRegistry，使用一个易读的名称 avlt.rotations。通过这种方法，我们可以轻松管理指标，客户端代码可以通过指标名称通过 MetricRegistry 检索任何指标。

实现很简单。我们使用 Python 字典来保存注册的指标。键是指标的名称，值是指标实例。我们还为支持的类型定义了 MetricType 以确保类型安全。

MetricType = Union[Counter, Histogram]
"""Alias for the supported metric types."""

class MetricRegistry:
    """A registry for metric instances."""

    def __init__(self) -> None:
        self._registry: Dict[str, MetricType] = dict()

    def register(self, name: str, metric: MetricType) -> None:
        """Given a metric, register it under the given name.

        Parameters
        ----------
        name: `str`
            The name of the metric

        metric: `MetricType`
            The type of the metric
        """
        self._registry[name] = metric

    def get_metric(self, name: str) -> MetricType:
        """Return the metric by the given name.

        Parameters
        ----------
        name: `str`
            The name of the metric

        Returns
        -------
        `MetricType`
            The metric instance by the given name.
        """
        return self._registry[name]

计数器

Counter 是一个简单的指标，用于计数某事，即增加和减少计数器。任何需要计数的事物都可以使用 Counter——例如，旋转次数。

class Counter:
    """An incrementing and decrementing counter metric."""

    def __init__(self) -> None:
        self._count = 0

    def increase(self, n: int = 1) -> None:
        """Increment the counter.

        Parameters
        ----------
        n: `int`
            The count to be increased.
        """
        self._count += n

    def decrease(self, n: int = 1) -> None:
        """Decrement the counter.

        Parameters
        ----------
        n: `int`
            The count to be decreased.
        """
        self._count -= n

    @property
    def count(self) -> int:
        """Return the current count."""
        return self._count

直方图

根据定义，直方图是对数值数据分布的图形表示。它是对连续变量（定量变量）概率分布的估计。Metrics 库中的 Histogram 测量数据集的分布，包括最小值、平均值、最大值、标准差和百分位数。

由于我们假设要测量的软件可以容纳在内存中，我们可以使用 Python 列表来保存所有值，并使用 Numpy 来计算数据集的分布。

import numpy as np

class Histogram:
    """A metric which calculates the distribution of a value."""

    def __init__(self) -> None:
        self._values: List[int] = list()

    def update(self, value: int) -> None:
        """Add a recorded value.

        Parameters
        ----------
        value: `int`
            value to be updated
        """
        self._values.append(value)

    def report(self) -> Dict:
        """Return the histogram report."""
        array = np.array(self._values)
        return {
            "min": array.min(),
            "max": array.max(),
            "medium": np.median(array),
            "mean": array.mean(),
            "stdDev": array.std(),
            "percentile": {
                "75": np.percentile(array, 75),
                "95": np.percentile(array, 95),
                "99": np.percentile(array, 99),
            },
        }

(完整实现可在 metrics.py 获取)

测试

一如既往，我们应该尽可能为我们的代码编写单元测试。有关完整的单元测试，请参阅 test_metrics.py。

将 Metrics 放入树中

红黑树

实现 Metrics 库后，我们可以将指标放在红黑树类中。由于我们要测量旋转次数，因此我们为此使用 Counter。为了计算树高度的变化，我们可以使用 Histogram。此外，我们需要一个 MetricRegistry 来注册 Counter 和 Histogram。MetricRegistry 支持管理整个程序的指标，因此更高层应将 MetricRegistry 实例传递给 RBTree 类。为此，我们修改 RBTree 的 __init__ 函数，如下所示：

class RBTree:
    def __init__(self, registry: Optional[metrics.MetricRegistry] = None) -> None:
        self._NIL: Leaf = Leaf()
        self.root: Union[Node, Leaf] = self._NIL
        self._metrics_enabled = True if registry else False
        if self._metrics_enabled and registry:
            self._rotate_counter = metrics.Counter()
            self._height_histogram = metrics.Histogram()
            registry.register(name="rbt.rotate", metric=self._rotate_counter)
            registry.register(name="rbt.height", metric=self._height_histogram)

(完整的实现请参阅 red_black_tree.py)

需要注意的一点是，虽然通过 Metrics 库之类的方法深入了解软件行为至关重要，但这也会带来成本。它增加了代码的复杂性，并可能降低其性能，因为它需要做更多的事情。因此，我们将 registry 参数设置为可选。当我们不对代码进行测量时，可以禁用指标以避免潜在的副作用。

旋转计数器

顾名思义，计数器需要在旋转发生时递增。因此，我们在旋转函数的末尾调用 increase() 函数，如下所示：

def _left_rotate(self, node_x: Node) -> None:
    node_y = node_x.right  # Set node y
    if isinstance(node_y, Leaf):  # Node y cannot be a Leaf
        raise RuntimeError("Invalid left rotate")
    # Turn node y's subtree into node x's subtree
    node_x.right = node_y.left
    if isinstance(node_y.left, Node):
        node_y.left.parent = node_x
    node_y.parent = node_x.parent
    # If node's parent is a Leaf, node y becomes the new root.
    if isinstance(node_x.parent, Leaf):
        self.root = node_y
    # Otherwise, update node x's parent.
    elif node_x == node_x.parent.left:
        node_x.parent.left = node_y
    else:
        node_x.parent.right = node_y
    node_y.left = node_x
    node_x.parent = node_y

    if self._metrics_enabled:
        self._rotate_counter.increase()

def _right_rotate(self, node_x: Node) -> None:
    node_y = node_x.left  # Set node y
    if isinstance(node_y, Leaf):  # Node y cannot be a Leaf
        raise RuntimeError("Invalid right rotate")
    # Turn node y's subtree into node x's subtree
    node_x.left = node_y.right
    if isinstance(node_y.right, Node):
        node_y.right.parent = node_x
    node_y.parent = node_x.parent
    # If node's parent is a Leaf, node y becomes the new root.
    if isinstance(node_x.parent, Leaf):
        self.root = node_y
    # Otherwise, update node x's parent.
    elif node_x == node_x.parent.right:
        node_x.parent.right = node_y
    else:
        node_x.parent.left = node_y
    node_y.right = node_x
    node_x.parent = node_y

    if self._metrics_enabled:
        self._rotate_counter.increase()

(完整的实现请参阅 red_black_tree.py)

高度直方图

执行插入和删除后，高度可能会发生变化，我们希望跟踪的是树高度的变化趋势。因此，我们可以在插入和删除的末尾更新 Histogram，如下所示：

def insert(self, key: Any, data: Any) -> None:
    # Color the new node as red.
    new_node = Node(key=key, data=data, color=Color.RED)
    parent: Union[Node, Leaf] = self._NIL
    current: Union[Node, Leaf] = self.root
    while isinstance(current, Node):  # Look for the insert location
        parent = current
        if new_node.key < current.key:
            current = current.left
        elif new_node.key > current.key:
            current = current.right
        else:
            raise tree_exceptions.DuplicateKeyError(key=new_node.key)
    new_node.parent = parent
    # If the parent is a Leaf, set the new node to be the root.
    if isinstance(parent, Leaf):
        new_node.color = Color.BLACK
        self.root = new_node
    else:
        if new_node.key < parent.key:
            parent.left = new_node
        else:
            parent.right = new_node
        # After the insertion, fix the broken red-black-tree-properties.
        self._insert_fixup(new_node)

    if self._metrics_enabled:
        self._height_histogram.update(value=self.get_height(self.root))

def delete(self, key: Any) -> None:
    if (deleting_node := self.search(key=key)) and (
        isinstance(deleting_node, Node)
    ):
        original_color = deleting_node.color
        # Case 1: no children or Case 2a: only one right child
        if isinstance(deleting_node.left, Leaf):
            replacing_node = deleting_node.right
            self._transplant(
                deleting_node=deleting_node, replacing_node=replacing_node
            )
            # Fixup
            if original_color == Color.BLACK:
                if isinstance(replacing_node, Node):
                    self._delete_fixup(fixing_node=replacing_node)
        # Case 2b: only one left child
        elif isinstance(deleting_node.right, Leaf):
            replacing_node = deleting_node.left
            self._transplant(
                deleting_node=deleting_node, replacing_node=replacing_node
            )
            # Fixup
            if original_color == Color.BLACK:
                if isinstance(replacing_node, Node):
                    self._delete_fixup(fixing_node=replacing_node)
        # Case 3: two children
        else:
            replacing_node = self.get_leftmost(deleting_node.right)
            original_color = replacing_node.color
            replacing_replacement = replacing_node.right
            # The replacing node is not the direct child of the deleting node
            if replacing_node.parent == deleting_node:
                if isinstance(replacing_replacement, Node):
                    replacing_replacement.parent = replacing_node
            else:
                self._transplant(replacing_node, replacing_node.right)
                replacing_node.right = deleting_node.right
                replacing_node.right.parent = replacing_node
            self._transplant(deleting_node, replacing_node)
            replacing_node.left = deleting_node.left
            replacing_node.left.parent = replacing_node
            replacing_node.color = deleting_node.color
            # Fixup
            if original_color == Color.BLACK:
                if isinstance(replacing_replacement, Node):
                    self._delete_fixup(fixing_node=replacing_replacement)

        if self._metrics_enabled:
            self._height_histogram.update(value=self.get_height(self.root))

(完整的实现请参阅 red_black_tree.py)

有了 Histogram，在进行一些插入和删除后，我们可以获得诸如最大值、平均值和中位数树高之类的报告。通过将这些数字与 Histogram 提供的标准差和百分位数结合起来，我们可以清晰地了解高度如何变化。例如，如果我们的 99% 百分位数接近中位数，并且标准差相对较低，我们可以说树高非常一致。换句话说，树相当平衡。

AVL 树

与红黑树类似，我们使用 Counter 来跟踪旋转，使用 Histogram 来监视树的高度。

class AVLTree:
    def __init__(self, registry: Optional[metrics.MetricRegistry] = None) -> None:
        self.root: Optional[Node] = None
        self._metrics_enabled = True if registry else False
        if self._metrics_enabled and registry:
            self._rotate_counter = metrics.Counter()
            self._height_histogram = metrics.Histogram()
            registry.register(name="avlt.rotate", metric=self._rotate_counter)
            registry.register(name="avlt.height", metric=self._height_histogram)

(完整的实现请参阅 avl_tree.py)

我们还在旋转函数的末尾增加 Counter，并在插入和删除后更新 Histogram。

def insert(self, key: Any, data: Any) -> None:
    new_node = Node(key=key, data=data)
    parent: Optional[Node] = None
    current: Optional[Node] = self.root
    while current:
        parent = current
        if new_node.key < current.key:
            current = current.left
        elif new_node.key > current.key:
            current = current.right
        else:
            raise tree_exceptions.DuplicateKeyError(key=new_node.key)
    new_node.parent = parent
    # If the tree is empty, set the new node to be the root.
    if parent is None:
       self.root = new_node
    else:
        if new_node.key < parent.key:
            parent.left = new_node
        else:
            parent.right = new_node
        if not (parent.left and parent.right):
            self._insert_fixup(new_node)

    if self._metrics_enabled:
        self._height_histogram.update(value=self.get_height(self.root))

def delete(self, key: Any) -> None:
    if self.root and (deleting_node := self.search(key=key)):
        # Case: no child
        if (deleting_node.left is None) and (deleting_node.right is None):
            self._delete_no_child(deleting_node=deleting_node)
        # Case: Two children
        elif deleting_node.left and deleting_node.right:
            replacing_node = self.get_leftmost(node=deleting_node.right)
            # Replace the deleting node with the replacing node,
            # but keep the replacing node in place.
            deleting_node.key = replacing_node.key
            deleting_node.data = replacing_node.data
            if replacing_node.right:  # The replacing node cannot have left child.
                self._delete_one_child(deleting_node=replacing_node)
            else:
                self._delete_no_child(deleting_node=replacing_node)
        # Case: one child
        else:
            self._delete_one_child(deleting_node=deleting_node)

        if self._metrics_enabled:
            self._height_histogram.update(value=self.get_height(self.root))

def _left_rotate(self, node_x: Node) -> None:
    node_y = node_x.right  # Set node y
    if node_y:
        # Turn node y's subtree into node x's subtree
        node_x.right = node_y.left
        if node_y.left:
            node_y.left.parent = node_x
        node_y.parent = node_x.parent
        # If node's parent is a Leaf, node y becomes the new root.
        if node_x.parent is None:
            self.root = node_y
        # Otherwise, update node x's parent.
        elif node_x == node_x.parent.left:
            node_x.parent.left = node_y
        else:
            node_x.parent.right = node_y
        node_y.left = node_x
        node_x.parent = node_y
        node_x.height = 1 + max(
            self.get_height(node_x.left), self.get_height(node_x.right)
        )
        node_y.height = 1 + max(
            self.get_height(node_y.left), self.get_height(node_y.right)
        )

        if self._metrics_enabled:
            self._rotate_counter.increase()

def _right_rotate(self, node_x: Node) -> None:
    node_y = node_x.left  # Set node y
    if node_y:
        # Turn node y's subtree into node x's subtree
        node_x.left = node_y.right
        if node_y.right:
            node_y.right.parent = node_x
        node_y.parent = node_x.parent
        # If node's parent is a Leaf, node y becomes the new root.
        if node_x.parent is None:
            self.root = node_y
        # Otherwise, update node x's parent.
        elif node_x == node_x.parent.right:
            node_x.parent.right = node_y
        else:
            node_x.parent.left = node_y
        node_y.right = node_x
        node_x.parent = node_y
        node_x.height = 1 + max(
            self.get_height(node_x.left), self.get_height(node_x.right)
        )
        node_y.height = 1 + max(
            self.get_height(node_y.left), self.get_height(node_y.right)
        )

        if self._metrics_enabled:
            self._rotate_counter.increase()

(完整的实现请参阅 avl_tree.py)

二叉搜索树

虽然本文重点关注 AVL 树和红黑树，但很容易将相同的方法应用于二叉搜索树。由于二叉搜索树没有旋转操作，我们只需要一个 Histogram 来测量树的高度。在下一节中，我们可以将二叉搜索树添加到比较中，以更好地了解自平衡树如何优于二叉搜索树。

(请参阅 binary_search_tree.py 以通过为二叉搜索树添加指标的实现)

使用 Metrics 比较红黑树和 AVL 树

实验 1

在第一个实验中，我们使用 Python 内置的随机抽样函数 ( random.sample) 生成一个包含 1000 个唯一随机整数的列表 ( insert_data)，这些整数是从 1 到 2000 的范围内选择的。并使用相同的 random.sample 函数将 insert_data 随机化为 delete_data。因此，从树中删除节点的顺序将是随机的。

import random

insert_data = random.sample(range(1, 2000), 1000)
delete_data = random.sample(insert_data, 1000)

然后，我们为指标定义了一个 MetricRegistry，并在定义树时传递它。在这里，我们还定义了一个二叉搜索树，以便我们可以看到二叉搜索树和自平衡树之间的性能差异。

from forest import metrics
from forest.binary_trees import avl_tree
from forest.binary_trees import binary_search_tree
from forest.binary_trees import red_black_tree

registry = metrics.MetricRegistry()
bstree = binary_search_tree.BinarySearchTree(registry=registry)
avltree = avl_tree.AVLTree(registry=registry)
rbtree = red_black_tree.RBTree(registry=registry)

一旦树和指标准备就绪，我们首先将 insert_data 插入树中，然后按照 delete_data 的顺序删除树中的节点。

for key in insert_data:
    bstree.insert(key=key, data=str(key))
    avltree.insert(key=key, data=str(key))
    rbtree.insert(key=key, data=str(key))

for key in delete_data:
    bstree.delete(key=key)
    avltree.delete(key=key)
    rbtree.delete(key=key)

最后，我们可以通过它们的名称检索指标，并打印出旋转次数和高度的直方图。

print("Binary Search Tree:")
bst_report = registry.get_metric(name="bst.height").report()
print(f"  Height:   {bst_report}")

print("AVL Tree:")
avlt_rotation_count = registry.get_metric(name="avlt.rotate").count
print(f"  Rotation: {avlt_rotation_count}")
avlt_report = registry.get_metric(name="avlt.height").report()
print(f"  Height:   {avlt_report}")

print("Red-Black Tree")
rbt_rotation_count = registry.get_metric(name="rbt.rotate").count
print(f"  Rotation: {rbt_rotation_count}")
rbt_repot = registry.get_metric(name="rbt.height").report()
print(f"  Height:   {rbt_repot}")

(完整代码可在 avlt_vs_rbt.py 获取)

由于输入数据是随机生成的，每次的结果可能会略有不同，但结果应该非常接近。输出可能如下所示：

Binary Search Tree:
  Height:   {'min': 0, 'max': 21, 'medium': 17.0, 'mean': 16.52976488244122, 'stdDev': 3.621953722665703, 'percentile': {'75': 20.0, '95': 20.0, '99': 21.0}}
AVL Tree:
  Rotation: 1053
  Height:   {'min': -1, 'max': 11, 'medium': 10.0, 'mean': 9.1765, 'stdDev': 1.7370514528936674, 'percentile': {'75': 10.0, '95': 11.0, '99': 11.0}}
Red-Black Tree
  Rotation: 647
  Height:   {'min': 0, 'max': 12, 'medium': 11.0, 'mean': 9.9605, 'stdDev': 1.78295814589126, 'percentile': {'75': 11.0, '95': 12.0, '99': 12.0}}

从输出中，我们可以看到 AVL 树和红黑树都非常平衡——AVL 树的最大高度为 11，红黑树为 12，中位数和平均值也与最大高度非常接近。此外，它们的百分位数相对于中位数，标准差也相对较低。因此，我们可以说两种树的高度都是一致的。

相比之下，二叉搜索树的高度更差——最大值为 21，中位数为 17。我们知道，根据我们在“二叉搜索树——分析”中提到的定理：在 n 个不重复键上随机构建的二叉搜索树的预期高度为 O(lg n)。即使我们随机生成二叉搜索树，其树高仍然比 AVL 树和红黑树差。此外，二叉搜索树的标准差比 AVL 树和红黑树更大。在这里，我们表明即使我们随机生成二叉搜索树，AVL 树和红黑树也比二叉搜索树更平衡。

关于旋转次数，结果表明 AVL 树比红黑树需要更多的旋转来保持平衡，这符合我们的预期。

实验 2

在第二个实验中，我们将比第一个实验更激进。我们随机生成插入和删除数据，并以随机顺序执行插入和删除。同样，我们使用 random.sample 函数生成一个包含 2000 个唯一随机整数的列表 ( insert_data)，这些整数是从 1 到 3000 的范围内选择的。但这次，我们生成一个包含 1000 个唯一随机整数的 delete_data，这些整数是从 1 到 3000 的范围内选择的。在这种情况下，插入和删除数据不相同。这意味着当我们尝试删除一个节点时，该节点可能不存在，但这不成问题。只要我们使用相同的 insert_data 和 delete_data 以相同的顺序测试我们的树，实验就是有效的。以下代码显示了我们如何生成测试数据。

import random

sample_data = random.sample(range(1, 3000), 2000)
insert_data = [("insert", item) for item in sample_data]

sample_data = random.sample(range(1, 3000), 1000)
delete_data = [("delete", item) for item in sample_data]

test_data = random.sample(
    (insert_data + delete_data), len(insert_data) + len(delete_data)
)

我们通过 random.sample 函数将 insert_data 和 delete_data 合并到一个 test_data 中，因此顺序将是随机的。此外，insert_data 和 delete_data 的每个项目都是一个包含操作类型（插入或删除）和要插入或删除的数据的对。因此，当我们遍历 test_data 时，我们知道需要执行哪些操作（插入或删除）。这就是我们确保所有三个树使用相同数据集和相同顺序的方法。以下代码演示了实验。

from forest import metrics
from forest.binary_trees import avl_tree
from forest.binary_trees import binary_search_tree
from forest.binary_trees import red_black_tree


registry = metrics.MetricRegistry()
bstree = binary_search_tree.BinarySearchTree(registry=registry)
avltree = avl_tree.AVLTree(registry=registry)
rbtree = red_black_tree.RBTree(registry=registry)

for operation, key in test_data:
    if operation == "insert":
        bstree.insert(key=key, data=str(key))
        avltree.insert(key=key, data=str(key))
        rbtree.insert(key=key, data=str(key))
    if operation == "delete":
        bstree.delete(key=key)
        avltree.delete(key=key)
        rbtree.delete(key=key)

print("Binary Search Tree:")
bst_report = registry.get_metric(name="bst.height").report()
print(f"  Height:   {bst_report}")

print("AVL Tree:")
avlt_rotation_count = registry.get_metric(name="avlt.rotate").count
print(f"  Rotation: {avlt_rotation_count}")
avlt_report = registry.get_metric(name="avlt.height").report()
print(f"  Height:   {avlt_report}")

print("Red-Black Tree")
rbt_rotation_count = registry.get_metric(name="rbt.rotate").count
print(f"  Rotation: {rbt_rotation_count}")
rbt_repot = registry.get_metric(name="rbt.height").report()
print(f"  Height:   {rbt_repot}")

(完整代码可在 avlt_vs_rbt_2.py 获取)

测试数据也是随机生成的，因此每次的结果可能略有不同。示例输出可能如下所示：

Binary Search Tree:
  Height:   {'min': 0, 'max': 23, 'medium': 22.0, 'mean': 20.394120153387302, 'stdDev': 3.249885374276778, 'percentile': {'75': 22.0, '95': 23.0, '99': 23.0}}
AVL Tree:
  Rotation: 1455
  Height:   {'min': 0, 'max': 12, 'medium': 11.0, 'mean': 10.175117170856412, 'stdDev': 1.6120322559944114, 'percentile': {'75': 11.0, '95': 12.0, '99': 12.0}}
Red-Black Tree
  Rotation: 1136
  Height:   {'min': 0, 'max': 13, 'medium': 11.0, 'mean': 10.826161056668086, 'stdDev': 1.8772598891928807, 'percentile': {'75': 12.0, '95': 13.0, '99': 13.0}}

我们得到了与实验 1 类似的结果——红黑树需要的旋转次数比 AVL 树少，AVL 树比红黑树稍微平衡一些，并且 AVL 树和红黑树都比二叉搜索树平衡得多。

摘要

AVL 树和红黑树在计算机软件中被广泛使用。它们各有优缺点。选择最适合我们用例的工具对于软件开发至关重要。此外，了解如何衡量我们的选择和实现对于软件的成功也很重要。本文介绍了使用我们在先前文章中实现的 AVL 树和红黑树作为示例来测量软件行为的方法。虽然本文示例的范围很小，但这里使用的概念可以应用于任何软件。希望本文能提供有价值的知识，让人们从中受益。