Vultr DocsLatest Content

Balance a Binary Search Tree

Updated on 19 May, 2025
Balance a Binary Search Tree header image

Problem Statement

In this problem, we are given the root of a binary search tree (BST) and our task is to transform it into a balanced BST using the same node values. A BST is deemed balanced if, for every node, the depths of its left and right subtrees differ by no more than one. The nature of this problem allows multiple valid solutions; therefore, any balanced BST that incorporates all the original node values correctly is an acceptable output.

Examples

Example 1

Input:

root = [1,null,2,null,3,null,4,null,null]

Output:

[2,1,3,null,null,null,4]
**Explanation:** This is not the only correct answer, [3,1,4,null,2] is also correct.

Example 2

Input:

root = [2,1,3]

Output:

[2,1,3]

Constraints

  • The number of nodes in the tree is in the range [1, 104].
  • 1 <= Node.val <= 105

Approach and Intuition

To convert an unbalanced BST into a balanced one, we can leverage the properties of in-order traversal and binary search methodology. Here's the step-by-step approach:

  1. In-Order Traversal:

    • Perform an in-order traversal of the given BST. This will allow us to retrieve all the node values in a sorted order (ascending).
    • Store these values in a list.

  2. Constructing a Balanced BST:

    • Utilize the sorted list of node values to construct a balanced BST. This can be efficiently done using a method similar to the binary search.
    • Define a recursive helper function that accepts the current segment of the sorted list (initially, the entire list).
    • Find the middle element of the current list segment - this element becomes the root of the current subtree. This choice ensures that the number of elements on either side of the root is as balanced as possible, contributing to a balanced depth of subtrees.
    • Recursively apply the same logic to the left half of the list to construct the left subtree and to the right half for the right subtree.

By following the above method, we ensure that the new BST is balanced by construction due to the even distribution of nodes on either side of each root chosen during the recursive calls. This process effectively minimizes the height of the tree, leading to a balanced BST where the depth of the two subtrees of every node never differs by more than 1.

This approach efficiently handles the constraints:

  • Given that the number of nodes can be up to 104, constructing the balanced BST in this manner ensures a time complexity that is logarithmic in nature relative to the number of nodes, keeping operations feasible within the provided limits.
  • All node values satisfy the given range requirement (1 <= Node.val <= 105), which are taken into consideration as they are directly used from the original tree without modification.

Solutions

  • C++
  • Java
  • Python
cpp 
class Solution {
public:
    TreeNode* balanceBST(TreeNode* root) {
        if (!root) return nullptr;

        TreeNode* dummyNode = new TreeNode(0);
        dummyNode->right = root;
        TreeNode* tempNode = dummyNode;
        while (tempNode->right) {
            if (tempNode->right->left) {
                rotateRight(tempNode, tempNode->right);
            } else {
                tempNode = tempNode->right;
            }
        }

        int countNodes = 0;
        tempNode = dummyNode->right;
        while (tempNode) {
            ++countNodes;
            tempNode = tempNode->right;
        }

        int m = pow(2, floor(log2(countNodes + 1))) - 1;
        performLeftRotations(dummyNode, countNodes - m);
        while (m > 1) {
            m /= 2;
            performLeftRotations(dummyNode, m);
        }

        TreeNode* balancedTop = dummyNode->right;
        delete dummyNode;
        return balancedTop;
    }

private:
    void rotateRight(TreeNode* parent, TreeNode* node) {
        TreeNode* newTop = node->left;
        node->left = newTop->right;
        newTop->right = node;
        parent->right = newTop;
    }

    void rotateLeft(TreeNode* parent, TreeNode* node) {
        TreeNode* newTop = node->right;
        node->right = newTop->left;
        newTop->left = node;
        parent->right = newTop;
    }

    void performLeftRotations(TreeNode* headNode, int count) {
        TreeNode* tmpNode = headNode;
        for (int i = 0; i < count; ++i) {
            TreeNode* rotateNode = tmpNode->right;
            rotateLeft(tmpNode, rotateNode);
            tmpNode = tmpNode->right;
        }
    }
};

The given C++ solution demonstrates how to balance a Binary Search Tree (BST). The primary objective is to transform a BST that may have become unbalanced into one that adheres to the optimal conditions of a balanced BST, where the left and right subtrees of every node differ in height by no more than one.

  • Initialize condition checks and subtle transformations to adjust tree balance by using rotations—specifically right and left rotations.
  • A dummyNode is utilized as a temporary node to facilitate easier rotations without handling special cases around the root directly.
  • The function rotateRight handles the right rotation of a node which is useful when a left-heavy subtree needs balancing.
  • The function rotateLeft is analogous but operates when a right-heavy imbalance occurs.
  • The process involves transforming the tree to a right-skewed form using right rotations where necessary.
  • Count the total number of nodes to apply transformations based on the tree's height and size.
  • Utilize log2 and pow functions to calculate optimum rotations and their respective levels in the tree.
  • After reformation to a right-skewed tree, the performLeftRotations function iteratively applies left rotations, transforming the tree to be more balanced with each step.
  • The final stage is cleanup where nodes are set to their final connections and unnecessary temporary nodes are removed to prevent memory leaks.

The functions and logic combined here ensure that the tree is compacted into a height-balanced tree, thus achieving minimal height for a given set of nodes and improving the efficiency of operations like insertion, deletion, and search within the BST.

java 
class Solution {

    public TreeNode constructBalancedBST(TreeNode root) {
        if (root == null) return null;

        TreeNode dummy = new TreeNode(0);
        dummy.right = root;
        TreeNode temp = dummy;
        while (temp.right != null) {
            if (temp.right.left != null) {
                rightRotation(temp, temp.right);
            } else {
                temp = temp.right;
            }
        }

        int count = 0;
        temp = dummy.right;
        while (temp != null) {
            count++;
            temp = temp.right;
        }

        int m = (int) Math.pow(2, Math.floor(Math.log(count + 1) / Math.log(2))) - 1;
        balanceTree(dummy, count - m);
        while (m > 1) {
            m /= 2;
            balanceTree(dummy, m);
        }

        TreeNode newRoot = dummy.right;
        return newRoot;
    }

    private void rightRotation(TreeNode parent, TreeNode child) {
        TreeNode temp = child.left;
        child.left = temp.right;
        temp.right = child;
        parent.right = temp;
    }

    private void leftRotation(TreeNode parent, TreeNode child) {
        TreeNode temp = child.right;
        child.right = temp.left;
        temp.left = child;
        parent.right = temp;
    }

    private void balanceTree(TreeNode head, int rotations) {
        TreeNode pointer = head;
        for (int i = 0; i < rotations; i++) {
            TreeNode next = pointer.right;
            leftRotation(pointer, next);
            pointer = pointer.right;
        }
    }
}

In the provided Java solution, the goal is to balance a binary search tree (BST) using rotations. The approach involves transforming the BST into a right-leaning linked list, then balancing the list to obtain a more optimal tree structure. Here's a breakdown of the strategy implemented in the code:

  1. The method constructBalancedBST starts by checking if the given root is null. If it is, it returns null immediately. Otherwise, it continues to set up a dummy node linked to the root to facilitate rotations.

  2. The input tree is first manipulated into a right-skewed form using a sequence of right rotations. This is done in a while loop checking for left children, then rotating them to the right until the entire tree is a right-leaning formation.

  3. Once the tree is skewed, the program calculates how many nodes would ideally exist in the left half of a perfectly balanced BST. This is handled using a logarithmic calculation to determine the largest complete subtree's size (m).

  4. Starting from this calculated point, left rotations are employed to stepwise balance large segments of the linked tree, reducing the imbalance in each step. Rotations are performed using the leftRotation method modifying parent and child links to enforce the left rotation.

  5. The process uses recursive tree balancing where larger portions of the skewed tree are incrementally balanced until the entire structure forms a balanced BST.

  6. The tree balancing operations are controlled by a for loop within balanceTree, which conducts a specified number of rotations on consecutive sub-trees, progressively adjusting the structure.

  7. The final tree rooted at newRoot once all rotations and rebalancing have occurred, is returned as the new balanced BST root.

The code successfully converts an unbalanced binary search tree into a balanced one and highlights the effective use of tree manipulation techniques such as rotations.

python 
class Solution:
    def balanceBinarySearchTree(self, rootNode: TreeNode) -> TreeNode:
        if not rootNode:
            return None

        # Creating a linear array-like structure from the BST
        vine_root = TreeNode(0)
        vine_root.right = rootNode
        pointer = vine_root
        while pointer.right:
            if pointer.right.left:
                self.performRightRotation(pointer, pointer.right)
            else:
                pointer = pointer.right

        # Counting the nodes in the vine
        total_nodes = 0
        pointer = vine_root.right
        while pointer:
            total_nodes += 1
            pointer = pointer.right

        # Rotating the vine to form a balanced BST
        full_tree = (1 << (math.floor(math.log2(total_nodes + 1)))) - 1
        self.rotationsToBalance(vine_root, total_nodes - full_tree)
        while full_tree > 1:
            full_tree >>= 1
            self.rotationsToBalance(vine_root, full_tree)

        balanced_tree_root = vine_root.right
        # Clearing the temporary head node
        vine_root = None
        return balanced_tree_root

    # Conducting a right rotation
    def performRightRotation(self, parent: TreeNode, rotateNode: TreeNode):
        temp = rotateNode.left
        rotateNode.left = temp.right
        temp.right = rotateNode
        parent.right = temp

    # Conducting a left rotation
    def performLeftRotation(self, parent: TreeNode, rotateNode: TreeNode):
        temp = rotateNode.right
        rotateNode.right = temp.left
        temp.left = rotateNode
        parent.right = temp

    # Function for handling the rotations to rebalance the tree
    def rotationsToBalance(self, head_node: TreeNode, rotations_count: int):
        current_node = head_node
        for _ in range(rotations_count):
            next_node = current_node.right
            self.performLeftRotation(current_node, next_node)
            current_node = current_node.right

The given Python code defines a method to balance a Binary Search Tree (BST) using the Day-Stout-Warren (DSW) algorithm. The approach involves converting the BST into a linked list "vine," and then transforming this vine into a balanced BST. The code embodies the solution inside a class named Solution with methods tailored for necessary operations like right rotation and left rotation.

Here's an outline of the operations performed:

  • Transforming BST to a VIne: Firstly, the code transforms the BST into a straight-line (vine) using right rotations. This is achieved by iterating through the tree and applying the performRightRotation method wherever left children exist.

  • Counting Nodes: Post vine creation, it counts the total number of nodes which is crucial for the next step.

  • Formation of Balanced BST: To balance the vine, the code determines the number of perfect binary tree nodes possible (full_tree). Starting with the largest full binary subtree possible, the tree is formed by applying left rotations through the rotationsToBalance method, progressively decreasing the subtree size by half until fully balanced.

  • Rotation Methods: The performRightRotation and performLeftRotation help in repositioning the nodes. The right rotation shifts a node down to the right, lifting its left child up, while the left rotation moves the node down to the left, elevating its right child.

This method effectively balances an unbalanced BST, ensuring improved operation times for insertions, deletions, and searches by maintaining a logarithmic height relative to the number of nodes. The transformation to a vine and subsequent balanced tree construction is systematic and ensures a depth-optimized BST.

Comments

No comments yet.