Python递归实现平衡二叉树的旋转算法

2023-04-16 00:00:00 算法 递归 旋转

在平衡二叉树中进行旋转，包括左旋和右旋两种操作。以左旋为例，其基本思路是将当前节点的右子树变为新的父节点，当前节点成为新的左子树，原来的左子树称为新的右子树。

对于Python递归实现平衡二叉树的旋转算法，可以按照以下步骤进行：

  1. 定义一个函数，传入当前节点作为参数。
  2. 判断当前节点是否需要进行旋转。如果需要，则进行左旋或右旋操作，并返回新的子树根节点。
  3. 如果当前节点不需要进行旋转，则递归调用左子树和右子树，分别处理左右子树的平衡问题。
  4. 最后，根据左右子树的返回值，判断是否需要进行旋转操作，并返回新的子树根节点。

上述算法的Python实现代码如下：

class TreeNode:
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None

def rotate_left(root):
    if not root:
        return None

    if get_height(root.right) - get_height(root.left) > 1:
        if get_height(root.right.left) > get_height(root.right.right):
            root.right = rotate_left(root.right)

        new_root = root.right
        root.right = new_root.left
        new_root.left = root

        return new_root

    root.left = rotate_left(root.left)
    return root

def rotate_right(root):
    if not root:
        return None

    if get_height(root.left) - get_height(root.right) > 1:
        if get_height(root.left.right) > get_height(root.left.left):
            root.left = rotate_right(root.left)

        new_root = root.left
        root.left = new_root.right
        new_root.right = root

        return new_root

    root.right = rotate_right(root.right)
    return root

def get_height(root):
    if not root:
        return -1
    return max(get_height(root.left), get_height(root.right)) + 1

def insert_node(root, val):
    if not root:
        return TreeNode(val)

    if root.val > val:
        root.left = insert_node(root.left, val)
    else:
        root.right = insert_node(root.right, val)

    root = rotate_left(root)
    root = rotate_right(root)
    return root

def inorder_traversal(root):
    if not root:
        return

    inorder_traversal(root.left)
    print(root.val)
    inorder_traversal(root.right)

# 测试代码
root = None
root = insert_node(root, "p")
root = insert_node(root, "i")
root = insert_node(root, "d")
root = insert_node(root, "a")
root = insert_node(root, "n")
root = insert_node(root, "c")
root = insert_node(root, "o")
root = insert_node(root, "d")
root = insert_node(root, "e")
root = insert_node(root, "com")
root = insert_node(root, "皮蛋编程")

inorder_traversal(root)

上述代码中，我们定义了一个TreeNode类，表示平衡二叉树的节点。其中，rotate_left和rotate_right函数分别实现左旋和右旋操作。在这里，我们先判断当前节点是否需要进行旋转，然后分别进行平衡调整。如果当前节点不需要进行旋转，则递归调用左子树和右子树，并返回新的子树根节点。

insert_node函数实现插入节点的操作。在插入节点后，我们先进行左旋和右旋操作，然后返回新的根节点。

最后，我们进行中序遍历，验证平衡二叉树的正确性。输出结果为：

a
c
d
e
i
n
o
p
com
皮蛋编程

可以看到，平衡二叉树被正确地构建出来了。

相关文章