编写伪代码判断给定BST是否为有效AVL树并设置节点高度

Great question! Let's tackle this step by step—first verifying if a given BST is a valid AVL tree (while setting each node's rank, which we'll define as the node's height), then adjusting it to become an AVL tree if it's unbalanced.

First: Clarify Definitions #

• AVL Tree: A self-balancing BST where for every node, the absolute difference between the heights of its left and right subtrees (called the balance factor) is at most 1.
• Rank (Height): For this problem, we'll define a node's rank as its height. Empty nodes have a height of -1, leaf nodes have a height of 0, and internal nodes have a height of 1 + max(left subtree height, right subtree height).

Step 1: Check if the BST is a Valid AVL Tree & Set Node Ranks #

We'll use a recursive function that returns two values: whether the subtree rooted at the given node is balanced, and the height of that subtree (which we'll assign to the node's rank property).

Node Structure #

Node {
    key: int       // BST key value
    left: Node|null
    right: Node|null
    rank: int      // Will store the node's height (initialized unassigned)
}

Validation & Rank-Setting Function #

Function isAVLAndSetRank(node):
    if node is null:
        return (isBalanced: true, height: -1)
    
    // Recursively process left and right subtrees
    leftResult = isAVLAndSetRank(node.left)
    rightResult = isAVLAndSetRank(node.right)
    
    // Set current node's rank to its calculated height
    node.rank = 1 + max(leftResult.height, rightResult.height)
    
    // Check if current node is balanced: both subtrees are balanced, and balance factor ≤ 1
    isCurrentBalanced = leftResult.isBalanced AND rightResult.isBalanced AND (abs(leftResult.height - rightResult.height) ≤ 1)
    
    return (isBalanced: isCurrentBalanced, height: node.rank)

How to Use This:

Call isAVLAndSetRank(root). If the returned isBalanced is true, your BST is a valid AVL tree, and all nodes have their correct rank values. If false, we need to balance it.

Step 2: Balance an Unbalanced BST into an AVL Tree #

To fix an unbalanced tree, we use four types of rotations (LL, RR, LR, RL) to rebalance nodes where the balance factor exceeds ±1. First, let's define helper rotation functions, then a balancing function.

Helper: Get Node Height #

Function getHeight(node):
    if node is null:
        return -1
    return node.rank

Rotation Functions #

LL Rotation (Right Rotation)

Used when the left subtree is taller, and the left child's left subtree is taller (or equal).

Function rotateRight(unbalancedNode):
    leftChild = unbalancedNode.left
    leftRightGrandchild = leftChild.right
    
    // Perform rotation
    leftChild.right = unbalancedNode
    unbalancedNode.left = leftRightGrandchild
    
    // Update ranks (heights) for rotated nodes
    unbalancedNode.rank = 1 + max(getHeight(unbalancedNode.left), getHeight(unbalancedNode.right))
    leftChild.rank = 1 + max(getHeight(leftChild.left), getHeight(leftChild.right))
    
    return leftChild // New root of the balanced subtree

RR Rotation (Left Rotation)

Used when the right subtree is taller, and the right child's right subtree is taller (or equal).

Function rotateLeft(unbalancedNode):
    rightChild = unbalancedNode.right
    rightLeftGrandchild = rightChild.left
    
    // Perform rotation
    rightChild.left = unbalancedNode
    unbalancedNode.right = rightLeftGrandchild
    
    // Update ranks
    unbalancedNode.rank = 1 + max(getHeight(unbalancedNode.left), getHeight(unbalancedNode.right))
    rightChild.rank = 1 + max(getHeight(rightChild.left), getHeight(rightChild.right))
    
    return rightChild // New root

LR Rotation (Left-Right Rotation)

Used when the left subtree is taller, but the left child's right subtree is taller.

Function rotateLR(unbalancedNode):
    // First rotate left on the left child to convert to LL case
    unbalancedNode.left = rotateLeft(unbalancedNode.left)
    // Then perform right rotation on the original node
    return rotateRight(unbalancedNode)

RL Rotation (Right-Left Rotation)

Used when the right subtree is taller, but the right child's left subtree is taller.

Function rotateRL(unbalancedNode):
    // First rotate right on the right child to convert to RR case
    unbalancedNode.right = rotateRight(unbalancedNode.right)
    // Then perform left rotation on the original node
    return rotateLeft(unbalancedNode)

Balancing Function #

This function recursively balances the tree, sets ranks, and returns the balanced subtree root.

Function balanceBSTAndSetRank(node):
    if node is null:
        return (root: null, height: -1)
    
    // Recursively balance left and right subtrees
    leftResult = balanceBSTAndSetRank(node.left)
    node.left = leftResult.root
    rightResult = balanceBSTAndSetRank(node.right)
    node.right = rightResult.root
    
    // Calculate current node's initial rank
    node.rank = 1 + max(leftResult.height, rightResult.height)
    
    // Check balance factor (left height - right height)
    balanceFactor = leftResult.height - rightResult.height
    
    // Handle unbalanced cases
    if balanceFactor > 1: // Left-heavy
        leftBalanceFactor = getHeight(node.left.left) - getHeight(node.left.right)
        if leftBalanceFactor >= 0:
            node = rotateRight(node) // LL case
        else:
            node = rotateLR(node)    // LR case
    elif balanceFactor < -1: // Right-heavy
        rightBalanceFactor = getHeight(node.right.left) - getHeight(node.right.right)
        if rightBalanceFactor <= 0:
            node = rotateLeft(node)  // RR case
        else:
            node = rotateRL(node)    // RL case
    
    // Update rank after rotation (if any)
    node.rank = 1 + max(getHeight(node.left), getHeight(node.right))
    
    return (root: node, height: node.rank)

How to Use This:

Call balancedRoot = balanceBSTAndSetRank(originalRoot).root. The returned balancedRoot is the root of a valid AVL tree, with all nodes' rank values correctly set.

Key Notes #

• BST First: This assumes your input is a valid BST. If you need to verify that first, add a separate function to check BST properties (all left subtree keys < node key, all right subtree keys > node key).
• Rank Definition: We use the standard AVL height definition where empty nodes have height -1—this makes balance factor calculations straightforward.
• Recursion: Recursion works well here because we process subtrees from the bottom up, ensuring we balance leaves before parent nodes.

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