平衡二叉树

基本介绍

左旋转调整成平衡二叉树

右旋转调整成平衡二叉树

双旋转调整成平衡二叉树

上述三种旋转方式的代码实现

class Node:
    """
    创建 Node 节点
    """
    value: int = 0
    left = None
    right = None

    def __init__(self, value: int):
        self.value = value

    def height(self):
        """
        返回以当前节点为根节点的树的高度
        """
        # 如果左子树不为空，则遍历左子树，最终获得左子树的高度
        # 如果右子树不为空，则遍历右子树，最终获得右子树的高度
        # 返回左右子树两个高度中的较大的那个值
        # 加1是因为当前节点要算一层高度
        return max(self.left.height() if self.left else 0, self.right.height() if self.right else 0) + 1

    def left_height(self):
        """
        返回左子树的高度
        """
        if self.left:
            return self.left.height()
        return 0

    def right_height(self):
        """
        返回右子树的高度
        """
        if self.right:
            return self.right.height()
        return 0

    def left_rotate(self):
        """
        通过左旋转的方式调整二叉树为平衡二叉树
        """
        # 以当前节点的值创建一个新的节点
        new_node = Node(self.value)
        # 把新节点的左子树设置为当前节点的左子树
        new_node.left = self.left
        # 把新节点的右子树设置为当前节点的右子树的左子树
        new_node.right = self.right.left
        # 把当前节点的值替换为当前节点右子节点的值
        self.value = self.right.value
        # 把当前节点的右子树设置为当前节点右子树的右子树
        self.right = self.right.right
        # 把当前节点的左子树设置为新节点
        self.left = new_node

    def right_rotate(self):
        """
        通过右旋转的方式调整二叉树为平衡二叉树
        """
        # 以当前节点的值创建一个新的节点
        new_node = Node(self.value)
        # 把新节点的右子树设置为当前节点的右子树
        new_node.right = self.right
        # 把新节点的左子树设置为当前节点的左子树的右子树
        new_node.left = self.left.right
        # 把当前节点的值替换为当前节点左子节点的值
        self.value = self.left.value
        # 把当前节点的左子树设置为当前节点左子树的左子树
        self.left = self.left.left
        # 把当前节点的右子树设置为新节点
        self.right = new_node

    def add(self, node):
        """
        添加节点
        node 表示要添加的节点
        """
        if node is None:
            return
        # 判断要添加节点和当前子树根节点的大小
        if node.value < self.value:
            # 要添加节点小于当前子树根节点
            if self.left is None:
                # 如果当前子树左子节点为空，则直接将要添加的节点挂在左子节点上
                self.left = node
            else:
                # 否则，递归当前子树的左节点，找到要放置的位置
                self.left.add(node)
        else:  # 要添加节点大于等于当前子树根节点
            if self.right is None:
                # 如果当前子树右子节点为空，则直接将要添加的节点挂在右子节点上
                self.right = node
            else:
                # 否则，递归当前子树的右节点，找到要放置的位置
                self.right.add(node)

        # 添加完一个节点后，如果 （右子树高度 - 左子树高度） > 1，则进行左旋转
        # 之所以叫左旋转，是因为右子树的高度高于左子树，需要将右子树调整（个人理解）
        if self.right_height() - self.left_height() > 1:
            # 如果当前节点右子树的左子树高度大于当前节点右子树的右子树高度
            # 则进行双旋转操作（先右旋转再左旋转）
            if self.right and self.right.left_height() > self.right.height():
                # 先对当前节点的右节点进行右旋转
                self.right.right_rotate()
                # 然后对当前节点进行左旋转
                self.left_rotate()
            else:  # 否则，直接左旋转
                self.left_rotate()
            # ！！！这里必须节return，否则会继续往下判断，有可能会出错
            return

        # 添加完一个节点后，如果 （左子树高度 - 右子树高度） > 1，则进行右旋转
        # 之所以叫右旋转，是因为左子树的高度高于右子树，需要将左子树调整（个人理解）
        if self.left_height() - self.right_height() > 1:
            # 如果当前节点左子树的右子树高度大于当前节点左子树的左子树高度
            # 则进行双旋转操作（先左旋转再右旋转）
            if self.left and self.left.right_height() > self.left.left_height():
                # 先对当前节点的左节点进行左旋转
                self.left.left_rotate()
                # 然后对当前节点进行右旋转
                self.right_rotate()
            else:  # 否则，直接右旋转
                self.right_rotate()

    def infix_order(self):
        """
        中序遍历二叉树
        """
        if self.left:
            self.left.infix_order()
        print(self.value, end=' ')
        if self.right:
            self.right.infix_order()


class AVLTree:
    """
    二叉排序树
    """
    root: Node = None

    def add(self, node: Node):
        """
        添加节点
        node: 要添加的节点
        """
        if self.root is None:
            # 如果根节点为空，直接将新节点当做根节点
            self.root = node
        else:
            # 否则，将新节点放在根节点下
            self.root.add(node)

    def infix_order(self):
        """
        中序遍历
        """
        if self.root:
            self.root.infix_order()
            print()
        else:
            print("二叉排序树为空...")


# 测试左旋转
arr = [4, 3, 6, 5, 7, 8]
avl_tree = AVLTree()
for i in arr:
    node = Node(i)
    avl_tree.add(node)
avl_tree.infix_order()
print("左旋转处理过后：")
print("树的高度：", avl_tree.root.height())
print("树的左子树高度：", avl_tree.root.left.height())
print("树的右子树高度：", avl_tree.root.right.height())
# 测试右旋转
arr = [10, 12, 8, 9, 7, 6]
avl_tree = AVLTree()
for i in arr:
    node = Node(i)
    avl_tree.add(node)
avl_tree.infix_order()
print("右旋转处理过后：")
print("树的高度：", avl_tree.root.height())
print("树的左子树高度：", avl_tree.root.left.height())
print("树的右子树高度：", avl_tree.root.right.height())

# 测试双旋转
arr = [10, 11, 7, 6, 8, 9]
avl_tree = AVLTree()
for i in arr:
    node = Node(i)
    avl_tree.add(node)
avl_tree.infix_order()
print("双旋转处理过后：")
print("树的高度：", avl_tree.root.height())
print("树的左子树高度：", avl_tree.root.left.height())
print("树的右子树高度：", avl_tree.root.right.height())