Scale Engineer

Balance a Binary Search Tree

Med
#1281Time: O(N log N), where N is the number of nodes. Traversing the original tree is O(N). We then perform N insertions into a tree that grows up to size N. Each insertion into a self-balancing tree takes O(log k) time, where k is the current size of the tree. The total time for N insertions is Σ(log k) for k=1 to N, which is O(N log N).Space: O(N). We need O(N) space to store the node values in a list. Additionally, the new tree itself requires O(N) space for its nodes. The recursion stack for insertions will take O(log N) space.
Patterns
Greedy
Algorithms
Divide and ConquerDepth-First Search
Data structures
TreeBinary TreeBinary Search Tree

Prompt

Given the root of a binary search tree, return a balanced binary search tree with the same node values. If there is more than one answer, return any of them.

A binary search tree is balanced if the depth of the two subtrees of every node never differs by more than 1.

 

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

Approaches

3 approaches with complexity analysis and trade-offs.

This approach simulates the behavior of a self-balancing Binary Search Tree, such as an AVL tree or a Red-Black tree. The core idea is to extract all the node values from the given unbalanced tree and then insert them one by one into a new tree structure that maintains its balance automatically after each insertion. While conceptually straightforward if a self-balancing tree library is available, implementing it from scratch is complex and less efficient time-wise compared to other methods.

Algorithm

  • Create a list to store node values.
  • Traverse the input BST to populate the list with all node values.
  • Initialize a new root to null.
  • For each value in the list, insert it into the new tree using a self-balancing insertion algorithm (like AVL insert).
  • The AVL insertion performs a standard BST insert, then updates heights and performs rotations if necessary to maintain balance.
  • Return the root of the newly constructed balanced tree.

Walkthrough

The algorithm proceeds in two main phases:

  1. Node Value Extraction: Traverse the input tree using any standard traversal method (like pre-order, in-order, or level-order) to collect all node values into a list.
  2. Incremental Construction: Create a new, initially empty, BST. Iterate through the collected values and insert each one into this new tree. The key is that the insertion logic must include a balancing step. For instance, in an AVL tree, after a standard BST insertion, we would trace back up to the root, checking the balance factor at each node and performing single or double rotations if the tree becomes unbalanced.

This process ensures that after all N nodes are inserted, the resulting tree is guaranteed to be balanced. However, the cost of rebalancing at each step accumulates, leading to a higher overall time complexity.

// Definition for a binary tree node.class TreeNode {    int val;    TreeNode left;    TreeNode right;    int height; // Extra field for AVL tree    TreeNode(int val) { this.val = val; this.height = 1; }} class Solution {    // Helper to get height    private int height(TreeNode N) {        if (N == null) return 0;        return N.height;    }     // Helper to get balance factor    private int getBalance(TreeNode N) {        if (N == null) return 0;        return height(N.left) - height(N.right);    }     // Right rotate    private TreeNode rightRotate(TreeNode y) { /* ... AVL rotation logic ... */ return new_root; }     // Left rotate    private TreeNode leftRotate(TreeNode x) { /* ... AVL rotation logic ... */ return new_root; }     // AVL insert function    private TreeNode insert(TreeNode node, int val) {        // 1. Standard BST insertion        if (node == null) return new TreeNode(val);        if (val < node.val) node.left = insert(node.left, val);        else if (val > node.val) node.right = insert(node.right, val);        else return node; // Duplicate values are not inserted         // 2. Update height of this ancestor node        node.height = 1 + Math.max(height(node.left), height(node.right));         // 3. Get the balance factor to check for unbalance        int balance = getBalance(node);         // 4. If unbalanced, perform rotations        // Left Left Case        if (balance > 1 && val < node.left.val) return rightRotate(node);        // Right Right Case        if (balance < -1 && val > node.right.val) return leftRotate(node);        // Left Right Case        if (balance > 1 && val > node.left.val) { /* ... */ }        // Right Left Case        if (balance < -1 && val < node.right.val) { /* ... */ }         return node;    }     private void inorderTraversal(TreeNode root, List<Integer> nodes) {        if (root == null) return;        inorderTraversal(root.left, nodes);        nodes.add(root.val);        inorderTraversal(root.right, nodes);    }     public TreeNode balanceBST(TreeNode root) {        List<Integer> nodes = new ArrayList<>();        inorderTraversal(root, nodes); // Get all nodes         TreeNode newRoot = null;        for (int val : nodes) {            newRoot = insert(newRoot, val); // Insert into AVL tree        }        return newRoot;    }}

Complexity

Time

O(N log N), where N is the number of nodes. Traversing the original tree is O(N). We then perform N insertions into a tree that grows up to size N. Each insertion into a self-balancing tree takes O(log k) time, where k is the current size of the tree. The total time for N insertions is Σ(log k) for k=1 to N, which is O(N log N).

Space

O(N). We need O(N) space to store the node values in a list. Additionally, the new tree itself requires O(N) space for its nodes. The recursion stack for insertions will take O(log N) space.

Trade-offs

Pros

  • The logic is straightforward if one has access to a pre-built self-balancing tree data structure.

Cons

  • Worse time complexity (O(N log N)) compared to the optimal O(N) solution.

  • Implementing a self-balancing tree from scratch is complex, verbose, and prone to errors.

Solutions

/** * 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 { private List < Integer > nums = new ArrayList <>(); public TreeNode balanceBST ( TreeNode root ) { dfs ( root ); return build ( 0 , nums . size () - 1 ); } private void dfs ( TreeNode root ) { if ( root == null ) { return ; } dfs ( root . left ); nums . add ( root . val ); dfs ( root . right ); } private TreeNode build ( int i , int j ) { if ( i > j ) { return null ; } int mid = ( i + j ) >> 1 ; TreeNode left = build ( i , mid - 1 ); TreeNode right = build ( mid + 1 , j ); return new TreeNode ( nums . get ( mid ), left , right ); } }

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next →

Next problem