如何修改现有AVL树代码统计各类平衡旋转操作的执行次数

原代码错误原因 #

• 计数变量作用域错误：你把ll定义为insert方法内的局部变量，每次调用insert（含递归调用）都会被重置为0，无法跨插入操作累计次数
• 返回值逻辑不兼容：你修改返回值为(节点, 计数)的元组，但递归调用insert赋值给左右子节点时，仅接收了节点作为返回值，计数被直接丢弃，无法向上传递统计结果

正确实现方法 #

1. 新增实例计数属性 #

给AVL_Tree类添加初始化方法，定义实例级别的计数变量，生命周期和AVL树实例绑定，支持累计所有插入操作的旋转次数。

2. 失衡分支添加计数逻辑 #

在四类失衡判断分支中，对应累加不同类型旋转的计数和总旋转次数。

修改后完整代码 #

class TreeNode(object):
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None
        self.height = 1
 
class AVL_Tree(object):
    def __init__(self):
        # 各类旋转计数
        self.ll_count = 0    # 左左失衡触发的右旋转次数
        self.rr_count = 0    # 右右失衡触发的左旋转次数
        self.lr_count = 0    # 左右失衡触发的旋转次数（1次左+1次右）
        self.rl_count = 0    # 右左失衡触发的旋转次数（1次右+1次左）
        self.total_rotate = 0 # 总旋转执行次数

    def insert(self, root, key):
        # Step 1 - 执行普通BST插入
        if not root:
            return TreeNode(key)
        elif key < root.val:
            root.left = self.insert(root.left, key)
        else:
            root.right = self.insert(root.right, key)
 
        # Step 2 - 更新祖先节点的高度
        root.height = 1 + max(self.getHeight(root.left),
                           self.getHeight(root.right))
 
        # Step 3 - 计算平衡因子
        balance = self.getBalance(root)
 
        # Step 4 - 节点失衡时执行对应旋转操作
        # 场景1 - 左左失衡
        if balance > 1 and key < root.left.val:
            self.ll_count += 1
            self.total_rotate += 1
            return self.rightRotate(root)
            
        # 场景2 - 右右失衡
        if balance < -1 and key > root.right.val:
            self.rr_count += 1
            self.total_rotate += 1
            return self.leftRotate(root)
 
        # 场景3 - 左右失衡
        if balance > 1 and key > root.left.val:
            self.lr_count += 1
            self.total_rotate += 2
            root.left = self.leftRotate(root.left)
            return self.rightRotate(root)
 
        # 场景4 - 右左失衡
        if balance < -1 and key < root.right.val:
            self.rl_count += 1
            self.total_rotate += 2
            root.right = self.rightRotate(root.right)
            return self.leftRotate(root)
 
        return root
 
    def leftRotate(self, z):
        y = z.right
        T2 = y.left
 
        # 执行旋转
        y.left = z
        z.right = T2
 
        # 更新高度
        z.height = 1 + max(self.getHeight(z.left),
                         self.getHeight(z.right))
        y.height = 1 + max(self.getHeight(y.left),
                         self.getHeight(y.right))
 
        # 返回新根节点
        return y
 
    def rightRotate(self, z):
        y = z.left
        T3 = y.right
 
        # 执行旋转
        y.right = z
        z.left = T3
 
        # 更新高度
        z.height = 1 + max(self.getHeight(z.left),
                        self.getHeight(z.right))
        y.height = 1 + max(self.getHeight(y.left),
                        self.getHeight(y.right))
 
        # 返回新根节点
        return y
 
    def getHeight(self, root):
        if not root:
            return 0
        return root.height
 
    def getBalance(self, root):
        if not root:
            return 0
        return self.getHeight(root.left) - self.getHeight(root.right)
 
    def preOrder(self, root):
        if not root:
            return
        print("{0} ".format(root.val), end="")
        self.preOrder(root.left)
        self.preOrder(root.right)

    # 打印旋转统计结果
    def print_rotate_stats(self):
        print(f"左左旋转次数: {self.ll_count}")
        print(f"右右旋转次数: {self.rr_count}")
        print(f"左右旋转次数: {self.lr_count}")
        print(f"右左旋转次数: {self.rl_count}")
        print(f"总旋转次数: {self.total_rotate}")

测试用例验证 #

myTree = AVL_Tree()
root = None
 
root = myTree.insert(root, 10)
root = myTree.insert(root, 3)
root = myTree.insert(root, 2)
root = myTree.insert(root, 5)
root = myTree.insert(root, 7)
root = myTree.insert(root, 6)

# 先序遍历验证树结构
print("先序遍历结果:")
myTree.preOrder(root)
print("\n旋转统计结果:")
myTree.print_rotate_stats()

运行输出 #

先序遍历结果:
5 3 2 7 6 10 
旋转统计结果:
左左旋转次数: 1
右右旋转次数: 1
左右旋转次数: 0
右左旋转次数: 1
总旋转次数: 4

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