二叉搜索树与平衡树

第 8 章 · 数据结构与算法

一、二叉搜索树（BST）

二叉搜索树（Binary Search Tree，BST）是一棵二叉树，且满足：对任意节点，其左子树中所有节点的值 < 该节点值，右子树中所有节点的值 > 该节点值。通常约定同一集合内无重复键（若有重复，可扩展为「左 ≤ 根 < 右」或挂在同一节点的链表上）。这一性质带来一个直接后果：中序遍历 BST 会得到有序序列。

1.1 查找、插入、删除

查找：从根开始，若目标等于当前节点则找到；若小于当前值则进入左子树，否则进入右子树；若走到空则不存在。一次比较排除一整棵子树，理想情况下高度为 O(log n)，故查找 O(log n)。

插入：先按查找走到底（若已存在可依约定处理）；在「空位」处挂上新节点。插入后仍满足 BST 性质。

删除：分三种情况。（1）被删节点是叶：直接删。（2）只有一个孩子：用该孩子顶替被删节点。（3）有两个孩子：通常用后继（中序后继，即右子树中最左节点）或前驱（中序前驱，即左子树中最右节点）的值覆盖当前节点，再在子树中删除该后继/前驱（该节点必至多一个孩子），从而转化为情况 1 或 2。

二、BST 的 Python 实现

下面给出带类型标注和详细注释的 BST：节点类、查找、插入、删除（用后继覆盖），以及中序遍历。代码约定不存重复键，重复插入视为无操作。

from __future__ import annotations
from typing import Optional, List


class BSTNode:
    """BST 节点：键 key、左子 left、右子 right。"""

    def __init__(self, key: int) -> None:
        self.key: int = key
        self.left: Optional[BSTNode] = None
        self.right: Optional[BSTNode] = None


def bst_search(root: Optional[BSTNode], key: int) -> Optional[BSTNode]:
    """在 BST 中查找 key，返回对应节点或 None。"""
    if root is None or root.key == key:
        return root
    if key < root.key:
        return bst_search(root.left, key)
    return bst_search(root.right, key)


def bst_insert(root: Optional[BSTNode], key: int) -> BSTNode:
    """在 BST 中插入 key（不允许重复）。返回插入后的根。"""
    if root is None:
        return BSTNode(key)
    if key < root.key:
        root.left = bst_insert(root.left, key)
    elif key > root.key:
        root.right = bst_insert(root.right, key)
    # key == root.key：不插入，直接返回
    return root


def _bst_min_node(node: BSTNode) -> BSTNode:
    """返回以 node 为根的子树中的最小键节点（最左节点）。"""
    while node.left is not None:
        node = node.left
    return node


def bst_delete(root: Optional[BSTNode], key: int) -> Optional[BSTNode]:
    """在 BST 中删除键为 key 的节点。返回删除后的根。"""
    if root is None:
        return None
    if key < root.key:
        root.left = bst_delete(root.left, key)
        return root
    if key > root.key:
        root.right = bst_delete(root.right, key)
        return root
    # 找到待删节点 root
    if root.left is None:
        return root.right
    if root.right is None:
        return root.left
    # 有两个孩子：用后继（右子树最左）的值覆盖，再删后继
    successor = _bst_min_node(root.right)
    root.key = successor.key
    root.right = bst_delete(root.right, successor.key)
    return root


def bst_inorder(root: Optional[BSTNode]) -> List[int]:
    """中序遍历 BST，返回有序键列表。"""
    if root is None:
        return []
    return bst_inorder(root.left) + [root.key] + bst_inorder(root.right)

三、平衡树：AVL 与红黑树概览

为了让树高稳定在 O(log n)，需要在插入、删除后做「平衡」：通过旋转调整结构。两种经典平衡树：

• AVL 树：任意节点的左右子树高度差（平衡因子）的绝对值 ≤ 1。平衡因子 = 左高 - 右高。插入/删除后若失衡，通过单旋或双旋恢复。
• 红黑树：通过「颜色」与五条规则保证从根到任意叶的路径长度相差不超过 2 倍，从而近似平衡。插入/删除时通过变色与旋转维护性质。

二者都能保证查找、插入、删除在 O(log n)；AVL 更「严格平衡」、查找略优，红黑树旋转与变色次数通常更少、写多读少场景常用（如很多语言的标准库 map/set 底层）。

四、AVL 树

AVL 树（由 Adelson-Velsky 和 Landis 提出）是 BST，且对每个节点有：|平衡因子| ≤ 1，其中平衡因子 = 左子树高度 - 右子树高度。空树高度记为 -1，单节点高度 0。

4.1 平衡因子与失衡类型

插入或删除后，从受影响节点沿父指针向上，可能首次遇到平衡因子变为 ±2 的节点，该节点为「失衡点」。以该节点为根的子树的失衡可归纳为四类（设失衡节点为 A，较高的一侧子节点为 B）：

• LL：A 的左子 B 的左子树更高（B 的平衡因子与 A 同号或 B 更高在左）。对 A 做右单旋。
• RR：A 的右子 B 的右子树更高。对 A 做左单旋。
• LR：A 的左子 B 的右子树更高。先对 B 左旋变成 LL，再对 A 右旋（或理解为对 A 做「左-右双旋」）。
• RL：A 的右子 B 的左子树更高。先对 B 右旋变成 RR，再对 A 左旋（右-左双旋）。

4.2 AVL 插入与旋转的 Python 实现

节点增加 height 字段（或每次递归时计算高度）；插入后自底向上更新高度并检查平衡因子，若为 ±2 则按 LL/RR/LR/RL 执行对应旋转。下面给出带高度与平衡因子的 AVL 节点、高度更新、四种旋转及插入的完整逻辑。

from __future__ import annotations
from typing import Optional, List


class AVLNode:
    """AVL 节点：key、left、right、height（当前子树高度，空为 -1）。"""

    def __init__(self, key: int) -> None:
        self.key: int = key
        self.left: Optional[AVLNode] = None
        self.right: Optional[AVLNode] = None
        self.height: int = 0


def _height(node: Optional[AVLNode]) -> int:
    """空树高度 -1，否则返回 node.height。"""
    return -1 if node is None else node.height


def _update_height(node: AVLNode) -> None:
    """根据左右子高度更新 node 的高度。"""
    node.height = 1 + max(_height(node.left), _height(node.right))


def _balance_factor(node: AVLNode) -> int:
    """平衡因子 = 左高 - 右高。"""
    return _height(node.left) - _height(node.right)


def _rotate_right(root: AVLNode) -> AVLNode:
    """右单旋：当前根为 A，左子为 B；旋转后 B 为根，A 为 B 的右子。"""
    new_root = root.left
    root.left = new_root.right
    new_root.right = root
    _update_height(root)
    _update_height(new_root)
    return new_root


def _rotate_left(root: AVLNode) -> AVLNode:
    """左单旋：当前根为 A，右子为 B；旋转后 B 为根，A 为 B 的左子。"""
    new_root = root.right
    root.right = new_root.left
    new_root.left = root
    _update_height(root)
    _update_height(new_root)
    return new_root


def avl_insert(root: Optional[AVLNode], key: int) -> AVLNode:
    """AVL 插入 key，返回新根。插入后自底向上检查平衡并旋转。"""
    if root is None:
        return AVLNode(key)
    if key < root.key:
        root.left = avl_insert(root.left, key)
    elif key > root.key:
        root.right = avl_insert(root.right, key)
    else:
        return root

    _update_height(root)
    bf = _balance_factor(root)

    # LL：左子左高 → 右单旋
    if bf == 2 and _balance_factor(root.left) >= 0:
        return _rotate_right(root)
    # RR：右子右高 → 左单旋
    if bf == -2 and _balance_factor(root.right) <= 0:
        return _rotate_left(root)
    # LR：左子右高 → 先左旋左子，再右旋根
    if bf == 2 and _balance_factor(root.left) < 0:
        root.left = _rotate_left(root.left)
        return _rotate_right(root)
    # RL：右子左高 → 先右旋右子，再左旋根
    if bf == -2 and _balance_factor(root.right) > 0:
        root.right = _rotate_right(root.right)
        return _rotate_left(root)
    return root

五、红黑树

红黑树（Red-Black Tree）是 BST，每个节点带颜色（红或黑），并满足以下五条性质：

1. 每个节点是红色或黑色。
2. 根是黑色。
3. 叶节点（NIL，空指针）视为黑色。
4. 红色节点的两个子节点都是黑色（不能有连续红）。
5. 从任意节点到其每个 NIL 叶的路径上，黑色节点数量相同（黑高一致）。

由性质 4 和 5 可证：从根到叶的路径长度最多相差 2 倍（一条路径全黑，另一条路径黑红相间），故树高 O(log n)。

5.1 红黑树插入（思路）

新插入节点先着红色（这样不破坏黑高）。若父为黑，直接结束。若父为红，则违反「不连续红」，需要调整：通过变色与旋转（与 AVL 类似的左旋/右旋），使根为黑、无连续红、黑高不变。具体分 uncle 为红 / 为黑且三角/线形等情形，经典教材与算法导论有完整分类。下面给出带父指针与颜色的节点定义，以及「插入后修复」的框架代码（仅示意 uncle 为红时向上变色）。

from __future__ import annotations
from typing import Optional
from enum import Enum


class Color(Enum):
    RED = 1
    BLACK = 2


class RBNode:
    """红黑树节点：key、color、left、right、parent（便于回溯修复）。"""

    def __init__(self, key: int, color: Color = Color.RED) -> None:
        self.key: int = key
        self.color: Color = color
        self.left: Optional[RBNode] = None
        self.right: Optional[RBNode] = None
        self.parent: Optional[RBNode] = None


def rb_insert_fixup(root: RBNode, z: RBNode) -> RBNode:
    """
    插入后修复：z 为新插入的红节点。若父为红则需调整。
    仅示例：当 uncle 为红时，将父与 uncle 变黑、祖父变红，再递归处理祖父。
    实际还需处理 uncle 为黑时的旋转（与 AVL 类似的左/右旋 + 变色）。
    """
    while z.parent is not None and z.parent.color == Color.RED:
        # 此处应区分父是祖父的左子还是右子，找 uncle，再分 uncle 红/黑
        # 若 uncle 为红：父、uncle 变黑，祖父变红，z = 祖父
        # 若 uncle 为黑：通过旋转 + 变色使当前子树根为黑
        break  # 简化：仅作框架，完整实现见算法导论或标准库源码
    root.color = Color.BLACK
    return root

5.2 AVL 与红黑树对比

六、小结与选用场景

BST 通过「左小右大」实现有序集合上的查找、插入、删除；若输入随机，期望 O(log n)，但顺序插入会退化。AVL 用平衡因子 ±1 与四种旋转保持严格平衡，查找最优。红黑树 用颜色与五条规则保证近似平衡，插入/删除时旋转与变色次数有界，工程中更常见。Python 标准库没有内置平衡 BST，若需要可选用 sortedcontainers 等第三方库；C++ 的 std::map、Java 的 TreeMap 多为红黑树实现。

七、例题与扩展

以下例题中假设二叉树节点类型为 TreeNode，包含 val、left、right，与第 7 章约定一致。

from typing import Optional


def is_valid_bst(root: Optional[TreeNode]) -> bool:
    """
    中序思路：BST 中序遍历严格递增。递归时传递当前子树允许的 (下界, 上界)，
    若节点值不在开区间 (lo, hi) 内则无效。初始为 (-inf, +inf)。
    """

    def in_range(node: Optional[TreeNode], lo: float, hi: float) -> bool:
        if node is None:
            return True
        if not (lo < node.val < hi):
            return False
        return in_range(node.left, lo, node.val) and in_range(node.right, node.val, hi)

    return in_range(root, float("-inf"), float("inf"))

from typing import Optional


def kth_smallest(root: Optional[TreeNode], k: int) -> int:
    """
    中序遍历 BST 得到有序序列，第 k 个即为第 k 小。
    迭代中序：用栈模拟，每次「向左走到底再弹栈访问」，计数到 k 即返回。
    """
    stack: list = []
    node = root
    count = 0
    while stack or node is not None:
        while node is not None:
            stack.append(node)
            node = node.left
        node = stack.pop()
        count += 1
        if count == k:
            return node.val
        node = node.right
    return -1  # k 合法时不会走到

from typing import Optional, List


def sorted_array_to_bst(nums: List[int]) -> Optional[TreeNode]:
    """
    有序数组的中序即自身。取中间元素为根，左半为左子树、右半为右子树，递归构造。
    这样左右子树节点数最多差 1，自然高度平衡。
    """

    def build(l: int, r: int) -> Optional[TreeNode]:
        if l > r:
            return None
        mid: int = (l + r) // 2
        root = TreeNode(nums[mid])
        root.left = build(l, mid - 1)
        root.right = build(mid + 1, r)
        return root

    return build(0, len(nums) - 1)

from typing import Optional


class BSTIterator:
    """中序迭代器：栈中保存「待访问的根」，next 时左到底、弹、转右。"""

    def __init__(self, root: Optional[TreeNode]) -> None:
        self._stack: list = []
        self._push_left(root)

    def _push_left(self, node: Optional[TreeNode]) -> None:
        while node is not None:
            self._stack.append(node)
            node = node.left

    def next(self) -> int:
        cur = self._stack.pop()
        self._push_left(cur.right)
        return cur.val

    def has_next(self) -> bool:
        return len(self._stack) > 0