Code Implementation

# Definition for a binary tree node.
class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

class Solution:
    def balanceBST(self, root: TreeNode) -> TreeNode:
        # Step 1: Inorder traversal to get sorted values
        def inorder(node):
            if not node:
                return []
            return inorder(node.left) + [node.val] + inorder(node.right)
        
        sorted_vals = inorder(root)
        
        # Step 2: Build balanced BST from sorted values
        def build_balanced_bst(nums, left, right):
            if left > right:
                return None
            mid = (left + right) // 2
            node = TreeNode(nums[mid])
            node.left = build_balanced_bst(nums, left, mid - 1)
            node.right = build_balanced_bst(nums, mid + 1, right)
            return node
        
        return build_balanced_bst(sorted_vals, 0, len(sorted_vals) - 1)
      

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
public:
    void inorder(TreeNode* node, vector<int>& vals) {
        if (!node) return;
        inorder(node->left, vals);
        vals.push_back(node->val);
        inorder(node->right, vals);
    }
    
    TreeNode* buildBalancedBST(vector<int>& vals, int left, int right) {
        if (left > right) return nullptr;
        int mid = (left + right) / 2;
        TreeNode* node = new TreeNode(vals[mid]);
        node->left = buildBalancedBST(vals, left, mid - 1);
        node->right = buildBalancedBST(vals, mid + 1, right);
        return node;
    }
    
    TreeNode* balanceBST(TreeNode* root) {
        vector<int> vals;
        inorder(root, vals);
        return buildBalancedBST(vals, 0, vals.size() - 1);
    }
};
      

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode() {}
 *     TreeNode(int val) { this.val = val; }
 *     TreeNode(int val, TreeNode left, TreeNode right) {
 *         this.val = val;
 *         this.left = left;
 *         this.right = right;
 *     }
 * }
 */
class Solution {
    public TreeNode balanceBST(TreeNode root) {
        List<Integer> vals = new ArrayList<>();
        inorder(root, vals);
        return buildBalancedBST(vals, 0, vals.size() - 1);
    }
    
    private void inorder(TreeNode node, List<Integer> vals) {
        if (node == null) return;
        inorder(node.left, vals);
        vals.add(node.val);
        inorder(node.right, vals);
    }
    
    private TreeNode buildBalancedBST(List<Integer> vals, int left, int right) {
        if (left > right) return null;
        int mid = (left + right) / 2;
        TreeNode node = new TreeNode(vals.get(mid));
        node.left = buildBalancedBST(vals, left, mid - 1);
        node.right = buildBalancedBST(vals, mid + 1, right);
        return node;
    }
}
      

/**
 * Definition for a binary tree node.
 * function TreeNode(val, left, right) {
 *     this.val = (val===undefined ? 0 : val)
 *     this.left = (left===undefined ? null : left)
 *     this.right = (right===undefined ? null : right)
 * }
 */
/**
 * @param {TreeNode} root
 * @return {TreeNode}
 */
var balanceBST = function(root) {
    const vals = [];
    function inorder(node) {
        if (!node) return;
        inorder(node.left);
        vals.push(node.val);
        inorder(node.right);
    }
    inorder(root);
    
    function buildBalancedBST(nums, left, right) {
        if (left > right) return null;
        const mid = Math.floor((left + right) / 2);
        const node = new TreeNode(nums[mid]);
        node.left = buildBalancedBST(nums, left, mid - 1);
        node.right = buildBalancedBST(nums, mid + 1, right);
        return node;
    }
    
    return buildBalancedBST(vals, 0, vals.length - 1);
};
      

Problem Description

You are given the root of a binary search tree (BST). Your task is to return a balanced version of this BST. A balanced BST is defined as a tree where the depth of the two subtrees of every node never differs by more than one. You must use all the original nodes of the tree (do not create new values or reuse values), and there is always at least one valid solution.

• The input is the root node of a BST.
• Return the root of a balanced BST containing all original nodes.
• Do not reuse or omit any node values.
• There may be multiple valid balanced BSTs; return any one of them.

Thought Process

When tackling this problem, we first recognize that a binary search tree can become unbalanced if, for example, nodes are inserted in sorted order, resulting in a tree that resembles a linked list. An unbalanced BST can degrade search performance from O(log n) to O(n). Our goal is to restructure the tree so that it is balanced, while still preserving the BST property (left child < parent < right child) and using all the original nodes.

A brute-force approach might involve repeatedly rotating nodes to try to balance the tree, but this is complex and inefficient. Instead, we notice that an in-order traversal of a BST yields a sorted list of node values. If we can rebuild the tree from this sorted list, always picking the middle element as the root (and recursively for sublists), we will end up with a balanced BST.

This leads us to a two-step approach: first, extract the sorted values using in-order traversal; second, use these values to construct a balanced BST by recursively choosing the middle node as the root of each subtree.

Solution Approach

To solve the problem efficiently, we follow these steps:

1. In-order Traversal:
   • Traverse the BST in-order (left, root, right) to collect all node values in a sorted array.
   • This preserves the BST property and ensures we use all original nodes.
2. Build Balanced BST:
   • Use the sorted array to build a balanced BST.
   • Recursively select the middle element of the array (or subarray) as the root node.
   • The left half of the array becomes the left subtree, and the right half becomes the right subtree.
   • Repeat recursively for each subtree until the whole tree is built.

This method ensures that the tree is as balanced as possible, with a minimal height, and all original node values are reused exactly once.

Example Walkthrough

Let's consider an example where the input BST is unbalanced:

Input BST:

        1
         \
          2
           \
            3
             \
              4
    

1. Step 1: In-order Traversal
   • Traverse the tree in-order: left, root, right.
   • Resulting sorted array: [1, 2, 3, 4]
2. Step 2: Build Balanced BST
   • Select the middle element (index 1, value 2) as the new root.
   • Left subtree: [1] (index 0)
   • Right subtree: [3, 4] (indices 2, 3)
   • For the right subtree, pick 3 as root, 4 as right child.

   The balanced BST will look like:

              2
            /   \
           1     3
                  \
                   4
           

The final tree is balanced: each node's left and right subtrees differ in height by at most 1.

Time and Space Complexity

• Brute-force Approach:
  • Repeatedly rotating nodes or checking balance at every insertion would have O(n^2) time complexity and is not practical for large trees.
• Optimized Approach (as described above):
  • Time Complexity: O(n) — In-order traversal takes O(n), and building the balanced BST takes O(n) as each node is visited once.
  • Space Complexity: O(n) — We store the node values in a list/array, and the recursion stack can be O(log n) for a balanced tree, but O(n) overall for storing the values.

Summary

To balance a binary search tree, the most efficient method is to perform an in-order traversal to extract the sorted node values, then rebuild the tree by recursively choosing the middle element as the root. This guarantees a balanced BST and leverages the properties of BSTs and sorted arrays. The approach is both simple and optimal, running in linear time and space, and is easy to implement in any modern programming language.