在 Python 中构建森林系列：AVL 树
5.00/5 (6投票s)
使用 Python 构建 AVL 树。
引言
在红黑树的讨论之后，本文将实现另一种自平衡二叉搜索树的变体：AVL 树。
项目设置
遵循构建森林系列中其他文章相同的风格和假设，该实现假定 Python 3.9 或更高版本。本文将向我们的项目添加两个模块：avl_tree.py 用于 AVL 树实现，test_avl_tree.py 用于其单元测试。添加这两个文件后，我们的项目布局将如下所示：
forest-python
├── forest
│ ├── __init__.py
│ ├── binary_trees
│ │ ├── __init__.py
│ │ ├── avl_tree.py
│ │ ├── binary_search_tree.py
│ │ ├── double_threaded_binary_tree.py
│ │ ├── red_black_tree.py
│ │ ├── single_threaded_binary_trees.py
│ │ └── traversal.py
│ └── tree_exceptions.py
└── tests
├── __init__.py
├── conftest.py
├── test_avl_tree.py
├── test_binary_search_tree.py
├── test_double_threaded_binary_tree.py
├── test_red_black_tree.py
├── test_single_threaded_binary_trees.py
└── test_traversal.py
（完整代码可在forest-python获取。）
什么是 AVL 树？
AVL 树（以发明者 Adelson-Velsky 和 Landis 的名字命名）是一种自平衡二叉搜索树。除了二叉搜索树的性质外，AVL 树还维护 AVL 树的性质以保持其平衡。
- 对于 AVL 树中的每个节点，其左子树和右子树的高度最多相差一。
该性质也称为平衡因子，可以改写为以下公式：
如果一个节点的平衡因子 > 0，我们称之为左偏。如果一个节点的平衡因子 < 0，我们称之为右偏。如果一个节点的平衡因子 = 0，则称之为平衡。
（请注意，有些人将平衡因子定义为右子树的高度减去左子树的高度。在这种情况下，左偏表示节点的平衡因子 < 0，而右偏发生在节点的平衡因子 > 0 时。但是，无论使用哪种定义，AVL 树的概念都是相同的。）
典型的 AVL 树可以如下图所示：
在上图中，BF 表示节点的平衡因子。节点下方的数字是节点的高度。如果一个节点为空，即 None，则其高度为 -1。这样可以更轻松地计算节点的平衡因子。
构建 AVL 树
本节将介绍 AVL 树的实现，并讨论实现选择背后的思考。
节点
我们不每次需要时都计算节点的高度，而是将高度存储在每个节点中。因此，节点结构比二叉搜索树节点多一个字段。
存储高度可以节省计算时间，因此我们在检查平衡因子时无需每次都计算高度。然而，这也有代价——当 AVL 树被修改时，例如插入或删除节点，我们需要保持高度的更新。有关高度更新的更多详细信息将在插入和删除部分中介绍。
与其他的二叉树节点一样，我们使用 dataclass 来定义 AVL 树节点。
from dataclasses import dataclass
@dataclass
class Node:
"""AVL Tree node definition."""
key: Any
data: Any
left: Optional["Node"] = None
right: Optional["Node"] = None
parent: Optional["Node"] = None
height: int = 0
类概述
与构建森林项目中的其他二叉树类型一样，AVL 树类具有类似的功能。
class AVLTree:
def __init__(self) -> None:
self.root: Optional[Node] = None
def __repr__(self) -> str:
"""Provie the tree representation to visualize its layout."""
if self.root:
return (
f"{type(self)}, root={self.root}, "
f"tree_height={str(self.get_height(self.root))}"
)
return "empty tree"
def search(self, key: Any) -> Optional[Node]:
…
def insert(self, key: Any, data: Any) -> None:
…
def delete(self, key: Any) -> None:
…
@staticmethod
def get_leftmost(node: Node) -> Node:
…
@staticmethod
def get_rightmost(node: Node) -> Node:
…
@staticmethod
def get_successor(node: Node) -> Optional[Node]:
…
@staticmethod
def get_predecessor(node: Node) -> Optional[Node]:
…
@staticmethod
def get_height(node: Optional[Node]) -> int:
…
def _get_balance_factor(self, node: Optional[Node]) -> int:
…
def _left_rotate(self, node_x: Node) -> None:
…
def _right_rotate(self, node_x: Node) -> None:
…
def _insert_fixup(self, new_node: Node) -> None:
…
def _transplant(self, deleting_node: Node, replacing_node: Optional[Node]) -> None:
…
def _delete_no_child(self, deleting_node: Node) -> None:
…
def _delete_one_child(self, deleting_node: Node) -> None:
…
def _delete_fixup(self, fixing_node: Node) -> None:
…
我们可以精确地像普通二叉搜索树一样实现大多数 AVL 树函数，例如搜索和大多数辅助函数。我们还可以使用二叉树遍历中的遍历函数来遍历 AVL 树。
插入和删除是可能导致 AVL 树失衡的两个操作。因此，AVLTree 类具有有助于保持树平衡的方法，包括 _left_rotate()、_right_rotate()、_insert_fixup() 和 _delete_fixup()。由于这些辅助方法主要用于保持 AVL 树平衡，因此我们将它们定义为私有函数，对客户端代码透明。
旋转
在插入或删除后恢复违反的 AVL 树性质的方法是旋转（类似于红黑树：旋转）。
下图演示了 AVL 树性质可能被破坏的四种情况：左-左、左-右、右-左和右-右。
当我们处理不平衡的 AVL 树时，我们总是从最底层的非平衡节点开始（即，节点的平衡因子为BF > 1或BF < -1）。因此，例如，上图中的节点 23 是最底层的非平衡节点。然后检查其高度较高的子节点的平衡因子。如果最底层的非平衡节点是左偏（即 BF > 0），并且其高度较高的子节点是左偏，则为左-左情况（上图中最左侧的情况）。如果高度较高的子节点的平衡因子是右偏（即 BF < 0），则为左-右情况（上图中第二个从左侧开始的情况）。图中的其他情况（右-左和右-右）与左-左情况和左-右情况对称。
以下子节将展示旋转如何重新平衡失衡情况。
右旋转（左-左情况）
对于左-左情况，我们在最底层的非平衡节点（本例中为节点 23）上执行右旋转。旋转后，节点 17 成为节点 23 的父节点。此外，节点 17 和节点 23 的高度和平衡因子在旋转后都发生了变化。
左-右旋转（左-右情况）
如果最底层的非平衡节点是左偏，但其子节点（高度较高者）是右偏，我们首先对子节点（本例中为节点 17）执行左旋转，然后对最底层的非平衡节点（节点 23）执行右旋转。
请注意，只有涉及旋转的节点才会改变其平衡因子和高度。例如，左旋转后，节点 17 和节点 21 改变了高度和平衡因子（黄色高亮）。右旋转后，只有节点 21 和节点 23 改变了高度和平衡因子（绿色表示高度）。
右-左旋转（右-左情况）
右-左情况与左-右情况对称。因此，我们先执行右旋转，然后执行左旋转。
左旋转（右-右情况）
右-右情况与左-左情况对称。因此，我们可以通过执行左旋转来使其平衡。
摘要
下表总结了不平衡的情况及其解决方案。
Insert
AVL 树的插入有一个重要影响：更新高度。因此，当我们向 AVL 树插入节点时，新节点可能会改变树的高度，从而违反 AVL 树性质。发生这种情况时，我们执行特定的旋转来重新平衡树。
高度更新
在实现 insert 函数之前，我们需要了解插入如何改变高度。首先，请注意，新插入的节点在插入后必须成为叶节点。因此，新节点的高度必须为 0，其平衡因子也为 0。此外，如果新节点在插入前其父节点已有子节点，则父节点和整个树的高度保持不变：高度不变，AVL 树性质未被破坏。
其次，仅当新节点插入前其父节点没有子节点时，才会发生高度变化。在这种情况下，我们更新新节点的祖先高度一直到根（如下图所示），只有新节点的祖先才有高度更新。当高度改变时，意味着可能会发生潜在的 AVL 树性质违反。
（上图并未违反 AVL 树性质。但是，我们将在接下来的章节中处理 AVL 树在插入后变得不平衡的情况。）
插入算法是对常规二叉搜索树插入算法的修改。
- 以与二叉搜索树插入相同的方式插入高度为 0 的新节点：通过从根开始遍历树并沿途比较新节点的键与每个节点的键，找到要插入新节点的正确位置（即新节点的父节点）。
- 更新高度并检查是否发生违反的 AVL 树性质，方法是从新节点开始回溯到根。在回溯到根的过程中，如果需要，更新沿途每个节点的高度。如果我们找到一个不平衡的节点，则执行特定的旋转来平衡它。旋转后，插入完成。如果未找到不平衡节点，则在到达根并更新其高度后，插入完成。
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
# After the insertion, fix the broken AVL-tree-property.
# If the parent has two children after inserting the new node,
# it means the parent had one child before the insertion.
# In this case, neither AVL-tree property breaks nor
# heights update requires.
if not (parent.left and parent.right):
self._insert_fixup(new_node)
修复
正如旋转部分所提到的，有四种潜在的不平衡情况。我们执行特定的旋转来恢复 AVL 树性质。在修复 AVL 树性质后，AVL 树将变得平衡。无需继续跟踪祖先的平衡因子。以下子节描述了每种情况的修复方法。
左-左情况
在上图中，我们添加了节点 19。插入后，我们从节点 17 及其向上开始检查节点 19 的祖先的平衡因子。然后，我们发现节点 37 的平衡因子不平衡且为左偏。我们还需要检查其高度较高的子节点，以确定执行哪种旋转。在插入情况下，高度较高的子节点必须发生在包含新节点的路径上，即本例中的节点 23，因为插入后节点 23 的平衡因子为 1。我们确定这是一种左-左情况。
因此，我们在节点 37 上执行右旋转。右旋转后，节点 23 成为新的根，节点 37 成为节点 23 的右子节点。请注意，只有涉及旋转的节点的高度和平衡因子才会改变。因此，在这种情况下，只有节点 23 和节点 37 的高度和平衡因子发生了变化。
为什么旋转后无需继续检查祖先的平衡因子？
插入前，节点 37 的高度为 2，节点 23 的高度为 1。插入后，节点 37 的高度变为 3，节点 23 的高度变为 2。然后我们执行了旋转。旋转后，节点 23 占据了节点 37 的原始位置，节点 23 的高度变为 2，节点 37 的高度变为 1。因此，同一位置的高度保持不变，即插入前根（节点 37）的高度为 2；旋转后根（节点 23）的高度仍为 2。换句话说，如果旋转前节点 37 有父节点（例如 x），旋转后，x 的左子节点将是节点 23，但 x 的高度不受影响。因此，我们无需在旋转后继续检查 x 的祖先的高度。
左-右情况
通过检查最底层非平衡节点（节点 37）及其高度较高的子节点（节点 23）的平衡因子来识别出这是左-右情况后，我们首先对节点 23 执行左旋转，使其变成左-左情况。然后，我们对非平衡节点（节点 37）执行右旋转。之后，我们恢复了被破坏的 AVL 树性质。
右-左情况
此情况与左-右情况对称，因此我们对非平衡节点（节点 37）的右子节点（节点 43）执行右旋转，使其变为右-右情况。然后，我们对非平衡节点（节点 37）执行左旋转。之后，我们修复了被破坏的 AVL 树性质。
左旋转（右-右情况）
右-右情况与左-左情况对称，因此我们可以通过执行左旋转来恢复其 AVL 树性质。
在修复分析之后，我们可以实现 _insert_fixup 函数，如下所示。请注意，我们在遍历树向上行走时始终更新节点的高度，并在旋转之前进行。
def _insert_fixup(self, new_node: Node) -> None:
parent = new_node.parent
while parent:
parent.height = 1 + max(
self.get_height(parent.left), self.get_height(parent.right)
)
grandparent = parent.parent
# grandparent is unbalanced
if grandparent:
if self._get_balance_factor(grandparent) > 1:
# Case Left-Left
if self._get_balance_factor(parent) >= 0:
self._right_rotate(grandparent)
# Case Left-Right
elif self._get_balance_factor(parent) < 0:
self._left_rotate(parent)
self._right_rotate(grandparent)
# Since the fixup does not affect the ancestor of the unbalanced
# node, exit the loop to complete the fixup process.
break
elif self._get_balance_factor(grandparent) < -1:
# Case Right-Right
if self._get_balance_factor(parent) <= 0:
self._left_rotate(grandparent)
# Case Right-Left
elif self._get_balance_factor(parent) > 0:
self._right_rotate(parent)
self._left_rotate(grandparent)
# Since the fixup does not affect the ancestor of the unbalanced
# node, exit the loop to complete the fixup process.
break
parent = parent.parent
搜索
搜索函数与二叉搜索树：搜索相同。
删除
AVL 树删除的基本思想与常规二叉搜索树相似。要删除的节点有三种情况：无子节点、只有一个子节点和有两个子节点。我们也使用相同的 transplant 方法将以 deleting_node 为根的子树替换为以 replacing_node 为根的子树。
移植
transplant 方法与二叉搜索树：删除相同。
def _transplant(self, deleting_node: Node, replacing_node: Optional[Node]) -> None:
if deleting_node.parent is None:
self.root = replacing_node
elif deleting_node == deleting_node.parent.left:
deleting_node.parent.left = replacing_node
else:
deleting_node.parent.right = replacing_node
if replacing_node:
replacing_node.parent = deleting_node.parent
与插入类似，AVL 树删除可能会更新树的高度并可能违反 AVL 树性质，因此我们需要检查 AVL 树在删除后是否变得不平衡并进行修复。
整体删除过程类似于常规二叉搜索树，但进行了一些修改。
- 找到要删除的节点（
deleting_node）。 - 如果
deleting_node没有子节点，则使用transplant方法将deleting_node替换为None。然后执行修复操作。 - 如果
deleting_node只有一个子节点，则使用transplant方法将deleting_node替换为其唯一的子节点。然后执行修复操作。 - 如果
deleting_node有两个子节点，则找到deleting_node的后继节点作为replacing_node。然后，用replacing_node的键和数据替换deleting_node的键和数据，从而用replacing_node替换deleting_node，但保持其原始平衡因子和高度（即高度和平衡因子不变，意味着 AVL 树性质未被违反）。之后，删除replacing_node，这与情况 2（无子节点）或情况 3（只有一个子节点）相同。
为了使实现清晰，我们定义了 _delete_no_child 方法来处理要删除的节点没有子节点的情况，以及 _delete_one_child 方法来处理要删除的节点只有一个子节点的情况。如果节点有两个子节点被删除，我们可以相应地重用 _delete_no_child 和 _delete_one_child。此外，由于有两个子节点的节点使用 _delete_no_child 和 _delete_one_child 删除，因此只有这两个方法需要调用 _delete_fixup 函数来修复不平衡的节点。因此，我们可以如下实现 delete、_delete_no_child 和 _delete_one_child 函数：
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)
def _delete_no_child(self, deleting_node: Node) -> None:
parent = deleting_node.parent
self._transplant(deleting_node=deleting_node, replacing_node=None)
if parent:
self._delete_fixup(fixing_node=parent)
def _delete_one_child(self, deleting_node: Node) -> None:
parent = deleting_node.parent
replacing_node = (
deleting_node.right if deleting_node.right else deleting_node.left
)
self._transplant(deleting_node=deleting_node, replacing_node=replacing_node)
if parent:
self._delete_fixup(fixing_node=parent)
修复
我们知道删除操作可能会改变高度并违反 AVL 树性质。与插入一样，有四种潜在的不平衡情况。我们执行特定的旋转来恢复 AVL 树性质。我们也从最底层的非平衡节点开始修复过程。最底层的非平衡节点必须是删除节点的祖先之一。然而，与插入修复不同的是，不平衡的平衡因子可能会传播到我们进行旋转的节点上方。因此，在恢复了最底层的非平衡节点后，我们需要检查其父节点。如果其父节点变得不平衡，则进行修复。重复此过程，直到到达根，并且根也平衡。
无子节点
正如我们在旋转部分提到的，四种情况可能违反 AVL 树性质。下图展示了这四种情况以及在删除节点没有子节点时如何恢复其平衡因子。
一个子节点
以下两张图展示了四种情况。
非平衡节点为左偏。
非平衡节点为右偏。
两个子节点
与其他二叉树的删除一样，我们可以将两个子节点的情况视为两种子情况：后继节点是 deleting_node 的直接子节点，以及 replacing_node 是 deleting_node 的最左侧节点。在这两种情况中，我们都可以通过替换 deleting_node 为 replacing_node 但保持原始高度和平衡因子来将两个子节点的情况转换为无子节点或单子节点情况。这样做，我们不改变高度和平衡因子。之后，我们删除 replacing_node，它变成无子节点情况或单子节点情况。下图展示了两个子节点删除如何工作以及如何修复其不平衡。
后继节点是删除节点的直接子节点。
后继节点是删除节点的最左侧节点。
完整修复操作的实现如下。与插入类似，我们在遍历树向上行走时更新节点的高度。
def _delete_fixup(self, fixing_node: Node) -> None:
while fixing_node:
fixing_node.height = 1 + max(
self.get_height(fixing_node.left), self.get_height(fixing_node.right)
)
if self._get_balance_factor(fixing_node) > 1:
# Case Left-Left
if self._get_balance_factor(fixing_node.left) >= 0:
self._right_rotate(fixing_node)
# Case Left-Right
elif self._get_balance_factor(fixing_node.left) < 0:
# The fixing node's left child cannot be empty
self._left_rotate(fixing_node.left)
self._right_rotate(fixing_node)
elif self._get_balance_factor(fixing_node) < -1:
# Case Right-Right
if self._get_balance_factor(fixing_node.right) <= 0:
self._left_rotate(fixing_node)
# Case Right-Left
elif self._get_balance_factor(fixing_node.right) > 0:
# The fixing node's right child cannot be empty
self._right_rotate(fixing_node.right)
self._left_rotate(fixing_node)
fixing_node = fixing_node.parent
辅助函数
辅助函数，例如获取最左侧节点和获取节点的后继节点，与二叉搜索树：辅助函数相同，实现可在Github 仓库中找到。与二叉搜索树不同的唯一辅助函数是获取高度的函数。
获取高度
由于每个节点都存储了其高度，因此获取节点的高度变得非常直接：只需返回高度即可。
@staticmethod
def get_height(node: Optional[Node]) -> int:
if node:
return node.height
# None has height -1
return -1
遍历
尽管 AVL 树节点比普通二叉搜索树节点多一个字段（即 height），但我们仍然可以使用我们在二叉树遍历中实现的完全相同的实现来遍历 AVL 树。我们唯一需要进行的修改是添加 AVL 树作为支持的类型。
# Alisa for the supported node types. For type checking.
SupportedNode = Union[None, binary_search_tree.Node, avl_tree.Node]
SupportedTree = Union[binary_search_tree.BinarySearchTree, avl_tree.AVLTree]
"""Alisa for the supported tree types. For type checking."""
（有关完整的源代码，请参阅 traversal.py。）
支持的类型用于类型检查，正如我们在二叉树遍历：函数接口中讨论过的。
测试
一如既往，我们应该尽可能为代码编写单元测试。请参阅 test_avl_tree.py 以获取完整的单元测试。
分析
AVL 树是一种自平衡二叉搜索树，其高度为 O(lg n)，其中 n 是节点数（这可以通过利用斐波那契数来证明。有关更多详细信息，请参阅 AVL 树维基百科）。因此，AVL 树的时间复杂度可总结如下表：
示例
与红黑树一样，AVL 树因其自平衡能力而被广泛应用于软件程序中。例如，本节使用我们在此实现的 AVL 树来实现键值对的 Map。
from typing import Any, Optional
from forest.binary_trees import avl_tree
from forest.binary_trees import traversal
class Map:
"""Key-value Map implemented using AVL Tree."""
def __init__(self) -> None:
self._avlt = avl_tree.AVLTree()
def __setitem__(self, key: Any, value: Any) -> None:
"""Insert (key, value) item into the map."""
self._avlt.insert(key=key, data=value)
def __getitem__(self, key: Any) -> Optional[Any]:
"""Get the data by the given key."""
node = self._avlt.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._avlt.delete(key=key)
def __iter__(self) -> traversal.Pairs:
"""Iterate the data in the map."""
return traversal.inorder_traverse(tree=self._avlt)
@property
def empty(self) -> bool:
"""Return `True` if the map is empty; `False` otherwise."""
return self._avlt.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)
（完整示例可在 avlt_map.py 获取。）
摘要
与红黑树一样，AVL 树在插入和删除操作上引入了一些复杂性以保持其平衡，但 AVL 树的自平衡能力为基本操作提供了 O(lg n) 的时间复杂度，这比常规二叉搜索树性能更好。
（最初发布于 Shun's Vineyard，日期为 2021 年 6 月 6 日。）
历史
- 2021 年 6 月 7 日：初始版本