Problem Description

You are given the root of a binary search tree (BST). Your task is to transform this BST into a balanced binary search tree while preserving all the original node values. A binary search tree is considered balanced if for every node in the tree, the depth difference between its left and right subtrees is at most 1.

The problem asks you to:

1. Take an existing BST (which may be unbalanced)
2. Reorganize its structure to make it balanced
3. Keep all the same values from the original tree
4. Return the root of the newly balanced BST

If multiple balanced configurations are possible, you can return any valid one.

The solution approach leverages the fact that an in-order traversal of a BST yields values in sorted order. The algorithm first performs an in-order traversal using the dfs function to collect all node values in the nums array in ascending order. Then, it reconstructs a balanced BST using the build function, which recursively:

• Selects the middle element as the root node (using mid = (i + j) >> 1, which is equivalent to (i + j) // 2)
• Builds the left subtree from elements before the middle
• Builds the right subtree from elements after the middle
• Returns a new TreeNode with the middle value and the constructed left and right subtrees

This divide-and-conquer approach ensures that the resulting tree is balanced because at each level, we're splitting the remaining elements roughly in half, preventing any path from root to leaf from being significantly longer than others.

Flowchart Walkthrough

First, let's pin down the algorithm using the Flowchart. Here's a step-by-step walkthrough:

Is it a graph?

• Yes: A binary tree is a special type of graph where each node has at most two children and there are no cycles.

Is it a tree?

• Yes: The problem explicitly deals with a binary search tree, which is a tree data structure.

DFS

• Yes: We arrive at DFS (Depth-First Search) as the recommended approach.

Conclusion: The flowchart suggests using DFS for this problem. This aligns perfectly with the solution approach, which uses DFS in two ways:

1. An in-order DFS traversal (dfs function) to extract all values from the original BST in sorted order
2. A recursive construction approach (build function) which is also a form of DFS, building the tree from top to bottom by recursively creating subtrees

The DFS pattern is ideal here because:

• We need to visit every node in the original tree to collect values (in-order traversal)
• We need to recursively construct a new balanced tree structure from the collected values
• Both operations naturally follow a depth-first pattern where we fully explore one branch before moving to another

Intuition

The key insight comes from understanding two fundamental properties:

1. BST Property: An in-order traversal of a BST always gives us values in sorted order
2. Balance Property: A balanced tree has roughly equal numbers of nodes in its left and right subtrees

When we see an unbalanced BST, we realize that the values themselves are fine - it's just their arrangement that's problematic. Think of it like having a sorted list of numbers that we've arranged poorly in tree form.

The natural question becomes: if we could start fresh with these sorted values, how would we build the most balanced tree possible?

The answer is surprisingly simple: always pick the middle element as the root. Why? Because this ensures we have approximately equal numbers of elements on both sides, which is exactly what balance means.

Consider a sorted array [1, 2, 3, 4, 5, 6, 7]. If we pick 4 (the middle) as root, we get [1, 2, 3] on the left and [5, 6, 7] on the right - perfectly balanced. We then apply the same logic recursively to each subarray.

This leads to our two-step approach:

1. Extract: Use DFS in-order traversal to get all values in sorted order (essentially "flattening" the unbalanced tree)
2. Rebuild: Recursively construct a new tree by always choosing the middle element as the root

The beauty of this solution is that it guarantees optimal balance - at each level of recursion, we're dividing the remaining elements as evenly as possible. The bit shift operation (i + j) >> 1 is just an efficient way to find the middle index, equivalent to (i + j) // 2.

Learn more about Greedy, Tree, Depth-First Search, Binary Search Tree, Divide and Conquer and Binary Tree patterns.

Solution Approach

The implementation follows the two-phase approach identified in our intuition: extract values via in-order traversal, then rebuild a balanced tree.

Phase 1: Extract Values with In-Order DFS

The dfs function performs an in-order traversal of the original BST:

def dfs(root: TreeNode):
    if root is None:
        return
    dfs(root.left)     # Visit left subtree
    nums.append(root.val)  # Process current node
    dfs(root.right)    # Visit right subtree

This traversal pattern (left → root → right) guarantees that values are collected in ascending order in the nums array. For example, if our unbalanced BST looks like:

    1
     \
      2
       \
        3

The nums array becomes [1, 2, 3].

Phase 2: Build Balanced Tree

The build(i, j) function constructs a balanced BST from the sorted array within index range [i, j]:

def build(i: int, j: int) -> TreeNode:
    if i > j:
        return None
    mid = (i + j) >> 1
    left = build(i, mid - 1)
    right = build(mid + 1, j)
    return TreeNode(nums[mid], left, right)

The algorithm works as follows:

1. Base Case: If i > j, the range is empty, return None
2. Find Middle: Calculate mid = (i + j) >> 1 (bit shift for integer division)
3. Recursive Construction:
   • Build left subtree from range [i, mid-1]
   • Build right subtree from range [mid+1, j]
   • Create root node with value nums[mid] and attach the subtrees

Example Walkthrough

For nums = [1, 2, 3]:

• Initial call: build(0, 2), mid = 1
  • Root = TreeNode(2)
  • Left child: build(0, 0) → creates node with value 1
  • Right child: build(2, 2) → creates node with value 3

Result:

    2
   / \
  1   3

The time complexity is O(n) where n is the number of nodes, as we visit each node exactly twice (once during traversal, once during construction). The space complexity is O(n) for storing the sorted values in the nums array.

Example Walkthrough

Let's walk through a concrete example to illustrate the solution approach.

Given unbalanced BST:

        1
         \
          2
           \
            3
             \
              4
               \
                5

Step 1: In-order Traversal (Extract Values)

We perform DFS in-order traversal to collect all values:

• Visit node 1: go left (none), add 1 to nums, go right
• Visit node 2: go left (none), add 2 to nums, go right
• Visit node 3: go left (none), add 3 to nums, go right
• Visit node 4: go left (none), add 4 to nums, go right
• Visit node 5: go left (none), add 5 to nums, go right (none)

Result: nums = [1, 2, 3, 4, 5]

Step 2: Build Balanced Tree

Now we recursively build from nums = [1, 2, 3, 4, 5]:

build(0, 4): range includes indices 0-4

• mid = (0 + 4) >> 1 = 2
• Create root with nums[2] = 3
• Build left subtree: build(0, 1)
  • mid = (0 + 1) >> 1 = 0
  • Create node with nums[0] = 1
  • Left: build(0, -1) returns None
  • Right: build(1, 1)
    • mid = 1, create node with nums[1] = 2
    • Both children are None
• Build right subtree: build(3, 4)
  • mid = (3 + 4) >> 1 = 3
  • Create node with nums[3] = 4
  • Left: build(3, 2) returns None
  • Right: build(4, 4)
    • mid = 4, create node with nums[4] = 5
    • Both children are None

Final balanced BST:

        3
       / \
      1   4
       \   \
        2   5

The tree is now balanced: every node has a depth difference of at most 1 between its left and right subtrees. The maximum path length from root to any leaf is 2, compared to 4 in the original unbalanced tree.

Solution Implementation

Time and Space Complexity

Time Complexity: O(n)

The algorithm consists of two main phases:

1. In-order traversal (dfs function): This visits each node exactly once to collect all values in sorted order, taking O(n) time.
2. Tree reconstruction (build function): This creates a new balanced BST by recursively dividing the sorted array and creating nodes. Each of the n nodes is created exactly once, taking O(n) time.

Therefore, the total time complexity is O(n) + O(n) = O(n).

Space Complexity: O(n)

The space usage includes:

1. nums array: Stores all n node values, requiring O(n) space.
2. Recursion stack for dfs: In the worst case (skewed tree), the recursion depth can be O(n).
3. Recursion stack for build: Since we're building a balanced tree, the maximum recursion depth is O(log n).
4. New tree nodes: We create n new TreeNode objects, requiring O(n) space.

The dominant space complexity is O(n), determined by the nums array and the space for the new tree nodes.

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Modifying the Original Tree Instead of Creating New Nodes

A common mistake is trying to reuse the original tree nodes instead of creating new TreeNode instances. This can lead to circular references or corrupted tree structure.

Incorrect approach:

def build_balanced_bst(left_idx: int, right_idx: int) -> TreeNode:
    if left_idx > right_idx:
        return None
  
    mid_idx = (left_idx + right_idx) // 2
    # Wrong: Trying to modify existing nodes
    root.val = sorted_values[mid_idx]  # This modifies the original tree
    root.left = build_balanced_bst(left_idx, mid_idx - 1)
    root.right = build_balanced_bst(mid_idx + 1, right_idx)
    return root

Correct approach: Always create new TreeNode instances to ensure a clean tree structure:

return TreeNode(sorted_values[mid_idx], left_subtree, right_subtree)

2. Integer Overflow in Middle Index Calculation

While less common in Python, calculating the middle index as (left + right) // 2 can theoretically cause integer overflow in languages with fixed integer sizes.

Safer calculation:

mid_idx = left_idx + (right_idx - left_idx) // 2

3. Forgetting to Handle Empty Tree

Not checking if the input root is None before processing can cause the algorithm to fail.

Solution: Add an early return check:

def balanceBST(self, root: TreeNode) -> TreeNode:
    if root is None:
        return None
  
    # Continue with the algorithm...

4. Using Global Variables Incorrectly

Declaring sorted_values as a class-level variable instead of within the function scope can cause issues when the method is called multiple times on the same Solution instance.

Incorrect:

class Solution:
    sorted_values = []  # Class variable - persists between calls
  
    def balanceBST(self, root: TreeNode) -> TreeNode:
        # sorted_values might contain values from previous calls

Correct: Always initialize the array within the method scope:

def balanceBST(self, root: TreeNode) -> TreeNode:
    sorted_values = []  # Local variable - fresh for each call

5. Off-by-One Errors in Index Boundaries

Using incorrect index boundaries when calling the build function can result in missing nodes or index out of bounds errors.

Common mistakes:

# Wrong: Using length instead of length-1
return build_balanced_bst(0, len(sorted_values))  

# Wrong: Starting from 1 instead of 0
return build_balanced_bst(1, len(sorted_values) - 1)

Correct:

return build_balanced_bst(0, len(sorted_values) - 1)