Balance a Binary Search Tree

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

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;
        }
    }
};