如何在AVL树中实现O(1)复杂度的节点索引及动态更新

核心思路 #

放弃维护每个节点的绝对index属性（中序遍历位置），转而维护**子树节点数size**属性。size的维护逻辑和AVL树的height完全一致，能在插入、删除、旋转时以O(1)时间更新单个节点的size，整体操作保持AVL树O(logn)的复杂度，同时基于size仍能高效实现索引范围查询。

具体实现步骤 #

1. 扩展节点结构 #

给AVLTreeNode添加size属性，初始值为1（新节点自身就是一棵子树）：

class AVLTreeNode:
    def __init__(self, key, item):
        self.key = key
        self.item = item
        self.left = None
        self.right = None
        self.height = 1  # 原有高度属性
        self.size = 1    # 新增：当前子树的总节点数（含自身）

2. 实现get_size辅助函数 #

和你已有的get_height逻辑一致，处理空节点的边界情况：

def get_size(self, current: AVLTreeNode) -> int:
    if current is not None:
        return current.size
    return 0

3. 插入操作中同步更新size #

在插入回溯时，更新完height后立即更新size，逻辑和height更新对称：

def insert_aux(self, current: AVLTreeNode, key: K, item: I) -> AVLTreeNode:
    if current is None:
        current = AVLTreeNode(key, item)
        self.length += 1
    elif key < current.key:
        current.left = self.insert_aux(current.left, key, item)
    elif key > current.key:
        current.right = self.insert_aux(current.right, key, item)
    else:
        raise ValueError('Inserting duplicate item')
    
    # 先更新高度，再更新子树节点数
    current.height = max(self.get_height(current.left), self.get_height(current.right)) + 1
    current.size = self.get_size(current.left) + self.get_size(current.right) + 1
    
    current = self.rebalance(current)
    return current

4. 删除操作中同步更新size #

删除操作同样在回溯阶段更新size，和height更新顺序一致：

def delete_aux(self, current: AVLTreeNode, key: K) -> AVLTreeNode:
    if current is None:
        raise ValueError('Item not found')
    
    if key < current.key:
        current.left = self.delete_aux(current.left, key)
    elif key > current.key:
        current.right = self.delete_aux(current.right, key)
    else:
        # 处理三种删除场景：叶子节点、单孩子节点、双孩子节点
        if current.left is None:
            self.length -= 1
            return current.right
        elif current.right is None:
            self.length -= 1
            return current.left
        else:
            # 用右子树最小节点替换当前节点
            successor = self.find_min(current.right)
            current.key = successor.key
            current.item = successor.item
            current.right = self.delete_aux(current.right, successor.key)
    
    # 更新高度和子树节点数
    current.height = max(self.get_height(current.left), self.get_height(current.right)) + 1
    current.size = self.get_size(current.left) + self.get_size(current.right) + 1
    
    current = self.rebalance(current)
    return current

5. 旋转操作中同步更新size #

旋转会改变节点的父子关系，需要先调整指针，再依次更新原根节点和新根节点的height与size：

def rotate_right(self, current: AVLTreeNode) -> AVLTreeNode:
    left_child = current.left
    temp = left_child.right
    
    # 执行右旋操作
    left_child.right = current
    current.left = temp
    
    # 先更新原根节点的height和size，再更新新根节点
    current.height = max(self.get_height(current.left), self.get_height(current.right)) + 1
    current.size = self.get_size(current.left) + self.get_size(current.right) + 1
    
    left_child.height = max(self.get_height(left_child.left), self.get_height(left_child.right)) + 1
    left_child.size = self.get_size(left_child.left) + self.get_size(left_child.right) + 1
    
    return left_child

def rotate_left(self, current: AVLTreeNode) -> AVLTreeNode:
    right_child = current.right
    temp = right_child.left
    
    # 执行左旋操作
    right_child.left = current
    current.right = temp
    
    # 先更新原根节点的height和size，再更新新根节点
    current.height = max(self.get_height(current.left), self.get_height(current.right)) + 1
    current.size = self.get_size(current.left) + self.get_size(current.right) + 1
    
    right_child.height = max(self.get_height(right_child.left), self.get_height(right_child.right)) + 1
    right_child.size = self.get_size(right_child.left) + self.get_size(right_child.right) + 1
    
    return right_child

6. 基于size实现索引范围查询 #

利用size可以快速定位第k个节点，进而实现高效的范围查询（以下示例索引从0开始）：

def get_kth_node(self, current: AVLTreeNode, k: int) -> AVLTreeNode:
    if current is None:
        return None
    
    left_size = self.get_size(current.left)
    if k < left_size:
        return self.get_kth_node(current.left, k)
    elif k == left_size:
        return current
    else:
        return self.get_kth_node(current.right, k - left_size - 1)

# 查询索引[start, end]范围内的节点
def query_range(self, start: int, end: int) -> list:
    result = []
    def traverse(node, base_pos):
        if node is None:
            return
        left_size = self.get_size(node.left)
        node_pos = base_pos + left_size
        
        if node_pos >= start and node_pos <= end:
            traverse(node.left, base_pos)
            result.append(node.item)
            traverse(node.right, node_pos + 1)
        elif node_pos < start:
            traverse(node.right, node_pos + 1)
        else:
            traverse(node.left, base_pos)
    
    traverse(self.root, 0)
    return result

方案优势 #

• 高效维护：size的更新逻辑和height完全同步，插入、删除、旋转操作均保持O(logn)复杂度，彻底避免原方案全树遍历的O(n)瓶颈。
• 功能兼容：基于size的索引查询逻辑和原基于绝对index的查询功能完全等价，无需修改上层业务逻辑。

内容的提问来源于stack exchange，提问作者ShaquilleOatmeal