1382. Balance a Binary Search Tree


Problem Description

The problem presents us with a binary search tree (BST) and asks us to transform it into a balanced binary search tree (BBST). A BST is considered balanced if, for every node in the tree, the depth difference between its left and right subtrees is no more than 1. The solution must maintain the original BST's node values and is not required to be unique—any correctly balanced BST with the original values is considered a valid answer.

Intuition

To balance a BST, we want to ensure that we minimize the depth differences between the left and right subtrees of all the nodes. A well-known approach to creating a balanced BST is to first convert the BST into a sorted array and then to rebuild the BST from this array.

The intuition behind this is that a sorted array will have the smaller values towards the beginning and larger values towards the end. By consistently selecting the middle element of the array (or subarray) to be the root node of the (sub)tree, we can ensure that the number of nodes on the left and right of any node is as equal as possible, thus creating a balanced tree.

The entire process consists of two main phases:

  1. In-Order Traversal (dfs function): Perform an in-order traversal of the BST and store each node's value in an array. An in-order traversal of a BST will process all nodes in ascending order, which aligns with the order we'd want in the array to reconstruct a BBST.

  2. Building the Balanced BST (build function): Using the sorted array, recursively select the middle element as the root node, then build the left subtree from the elements before it and the right subtree from the elements after it. This step is akin to a binary search where we use the middle of the current array (or subarray) as the root and apply the logic recursively to the left and right halves.

The provided solution implements these two key phases with the help of two helper functions: dfs for in-order traversal and build for constructing the balanced BST from the sorted values array.

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

Solution Approach

The implementation of the solution is done in two parts as per the problem requirement: first, the BST is converted to a sorted list, and then the list is used to construct a balanced BST.

Conversion of BST to Sorted List

The conversion of the BST to a sorted list is done using an in-order traversal, which is a depth-first search (DFS). Before the traversal, an empty list vals is initialized. The dfs function is defined to perform the traversal:

  1. If the current node (root) is None, the function returns, as it means we've reached a leaf node's child.
  2. It recursively calls itself on the left child (dfs(root.left)), processing all left descendants first (which are smaller in a BST).
  3. The value of the current node is appended to the list vals (vals.append(root.val)).
  4. It recursively calls itself on the right child (dfs(root.right)), processing all right descendants afterward (which are larger in a BST).

After defining the dfs function, it is called initially with the root of the BST. This will result in vals containing all the node values in sorted order since BSTs inherently maintain an in-order sequence of ascending values.

Building the Balanced BST

The balanced BST is constructed using the build function which follows a divide and conquer strategy:

  1. The function takes two indices, i and j, which represent the start and end of the current subarray within vals we're using to build the subtrees.
  2. The base case checks if i is greater than j. If so, it returns None, as this means there are no more elements to create nodes from in this subarray.
  3. The middle index mid is calculated using (i + j) >> 1, which is a bitwise operation equivalent to dividing the sum by 2 but is faster to compute.
  4. A new TreeNode is created using vals[mid] as the root value.
  5. The left subtree of the root is built by recursively calling build(i, mid - 1), effectively using the left half of the current subarray.
  6. The right subtree of the root is built by recursively calling build(mid + 1, j), using the right half of the current subarray.

The build function initializes the process by being called with the start 0 and end len(vals) - 1 indices, essentially representing the entirety of the sorted list.

Finally, the balanceBST function returns the root node of the newly constructed balanced BST, which results from calling build(0, len(vals) - 1).

Throughout the solution, we make use of recursion and divide-and-conquer techniques to arrive at a balanced BST that meets the depth difference requirement for all nodes.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

What are the two properties the problem needs to have for dynamic programming to be applicable? (Select 2)

Example Walkthrough

Let's consider a binary search tree (BST) that is skewed to the right as our example:

1    1
2     \
3      2
4       \
5        3
6         \
7          4

This BST is not balanced because the right subtree of each node has a height that is more than 1 greater than the left subtree (which is non-existent). Our objective is to balance this BST according to the solution approach described above.

Conversion of BST to Sorted List

Following the in-order traversal, we obtain the values from the tree in sorted order. We start with an empty list vals.

  1. Begin at the root (1). Left child is None, so move to step 3.
  2. Append 1 to vals (thus vals = [1]).
  3. Move to the right child, repeat steps for node 2, and so on.

After traversing all the nodes, the vals list becomes:

1vals = [1, 2, 3, 4]

This list represents the in-order traversal of the given BST and is sorted.

Building the Balanced BST

Now, we use the sorted list vals to construct a balanced BST.

  1. Call build(0, 3) since vals has four elements, indices ranging from 0 to 3.
  2. Calculate mid = (0 + 3) >> 1, resulting in mid = 1.
  3. Create a new TreeNode with vals[1] (which is 2), making it the root.
  4. For the left subtree, we call build(0, 0):
    • Calculate mid = (0 + 0) >> 1, resulting in mid = 0.
    • Create a new TreeNode with vals[0] (which is 1).
    • No more elements to the left or right, so the left subtree is just the node with 1.
  5. For the right subtree, we call build(2, 3):
    • Calculate mid = (2 + 3) >> 1, resulting in mid = 2.
    • Create a new TreeNode with vals[2] (which is 3).
    • For the left subtree of 3, there's no element, so return None.
    • For the right subtree of 3, call build(3, 3):
      • Only one element left, so create a new TreeNode with vals[3] (which is 4).
      • No more elements to the left or right.

After completing these steps, we have successfully constructed the following balanced BST:

1     2
2    / \
3   1   3
4        \
5         4

This resulting tree is balanced as the depth difference between the left and right subtrees for all nodes is no more than 1. It maintains all the original values from the input BST and is a valid solution according to the problem description.

Solution Implementation

1# Definition for a binary tree node.
2class TreeNode:
3    def __init__(self, val=0, left=None, right=None):
4        self.val = val
5        self.left = left
6        self.right = right
7
8class Solution:
9    def balanceBST(self, root: TreeNode) -> TreeNode:
10        # Helper function to perform inorder traversal and collect values in a list.
11        def inorder_traversal(node):
12            if node is None:  # Base case: If node is None, return.
13                return
14            inorder_traversal(node.left)  # Recurse on left subtree.
15            node_values.append(node.val)  # Append the value of the current node.
16            inorder_traversal(node.right)  # Recurse on right subtree.
17
18        # Helper function to build a balanced BST given a sorted value list.
19        def build_balanced_bst(start, end):
20            if start > end:  # If starting index is greater than ending, subtree is empty.
21                return None
22          
23            mid = (start + end) // 2  # Compute the middle index. 
24            node = TreeNode(node_values[mid])  # Create a node with the middle value.
25          
26            # Recursively build left and right subtrees using the split lists.
27            node.left = build_balanced_bst(start, mid - 1)
28            node.right = build_balanced_bst(mid + 1, end)
29            return node  # Return the newly created node.
30
31        node_values = []  # Initialize an empty list to store the tree node values.
32        inorder_traversal(root)  # Fill the node_values list with values from the BST.
33        # Build and return a balanced BST using the list of values.
34        return build_balanced_bst(0, len(node_values) - 1)
35
1class Solution {
2    // Class member to hold the tree values in sorted order.
3    private List<Integer> treeValues;
4
5    // Method to turn a binary search tree into a balanced binary search tree.
6    public TreeNode balanceBST(TreeNode root) {
7        treeValues = new ArrayList<>();
8        // Populate the treeValues list with the values in sorted order.
9        inOrderTraversal(root);
10        // Reconstruct the tree in a balanced manner.
11        return buildBalancedTree(0, treeValues.size() - 1);
12    }
13
14    // Helper method to perform an in-order traversal of the tree and store the values.
15    private void inOrderTraversal(TreeNode node) {
16        if (node == null) {
17            // Base case: if the node is null, do nothing.
18            return;
19        }
20        // Traverse the left subtree first.
21        inOrderTraversal(node.left);
22        // Store the value of the current node.
23        treeValues.add(node.val);
24        // Traverse the right subtree next.
25        inOrderTraversal(node.right);
26    }
27
28    // Helper method to build a balanced binary search tree using the stored values.
29    private TreeNode buildBalancedTree(int start, int end) {
30        // Base case: if start index is greater than end index, subtree is null.
31        if (start > end) {
32            return null;
33        }
34        // Find the mid point to make the current root.
35        int mid = start + (end - start) / 2;
36        // Create a new TreeNode with the mid value.
37        TreeNode root = new TreeNode(treeValues.get(mid));
38        // Recursively build the left subtree.
39        root.left = buildBalancedTree(start, mid - 1);
40        // Recursively build the right subtree.
41        root.right = buildBalancedTree(mid + 1, end);
42        // Return the node.
43        return root;
44    }
45}
46
1// Definition for a binary tree node.
2struct TreeNode {
3    int val;
4    TreeNode *left;
5    TreeNode *right;
6    TreeNode() : val(0), left(nullptr), right(nullptr) {}
7    TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
8    TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
9};
10
11class Solution {
12public:
13    vector<int> nodeValues; // Vector to store the values of nodes in sorted order.
14
15    // Main function to create a balanced BST from the unbalanced BST `root`.
16    TreeNode* balanceBST(TreeNode* root) {
17        // Perform in-order traversal to get the values in the sorted order.
18        inOrderTraversal(root);
19        // Rebuild the tree using the sorted values from `nodeValues`.
20        return rebuildTree(0, nodeValues.size() - 1);
21    }
22
23    // Helper function to perform an in-order traversal of a BST.
24    void inOrderTraversal(TreeNode* node) {
25        if (!node) return; // Base case: if the node is null, do nothing.
26        inOrderTraversal(node->left); // Traverse the left subtree.
27        nodeValues.push_back(node->val); // Visit the node and add its value to `nodeValues`.
28        inOrderTraversal(node->right); // Traverse the right subtree.
29    }
30
31    // Helper function to build a balanced BST from the sorted values.
32    TreeNode* rebuildTree(int start, int end) {
33        if (start > end) return nullptr; // Base case: no nodes to construct the tree.
34
35        // Calculate the middle index and use it as the root to balance the tree.
36        int mid = (start + end) / 2;
37        TreeNode* newNode = new TreeNode(nodeValues[mid]); // Create a new node with the mid value.
38
39        // Recursively construct the left and right subtrees and link them to the new node.
40        newNode->left = rebuildTree(start, mid - 1);
41        newNode->right = rebuildTree(mid + 1, end);
42
43        return newNode; // Return the new subtree root.
44    }
45};
46
1// Definition for a binary tree node.
2class TreeNode {
3    val: number;
4    left: TreeNode | null;
5    right: TreeNode | null;
6
7    constructor(val: number = 0, left: TreeNode | null = null, right: TreeNode | null = null) {
8        this.val = val;
9        this.left = left;
10        this.right = right;
11    }
12}
13
14// Array to store the values of nodes in sorted order
15const nodeValues: number[] = [];
16
17// Main function to create a balanced BST from the unbalanced BST 'root'
18function balanceBST(root: TreeNode | null): TreeNode | null {
19    if (root === null) return null; // Guard clause for null input
20  
21    // Perform in-order traversal to get the values in sorted order
22    inOrderTraversal(root);
23    // Rebuild the tree using the sorted values from 'nodeValues'
24    return rebuildTree(0, nodeValues.length - 1);
25}
26
27// Helper function to perform an in-order traversal of a BST
28function inOrderTraversal(node: TreeNode | null): void {
29    if (node === null) return; // Base case: if the node is null, do nothing
30    inOrderTraversal(node.left); // Traverse the left subtree
31    nodeValues.push(node.val); // Visit the node and add its value to 'nodeValues'
32    inOrderTraversal(node.right); // Traverse the right subtree
33}
34
35// Helper function to build a balanced BST from the sorted values
36function rebuildTree(start: number, end: number): TreeNode | null {
37    if (start > end) return null; // Base case: no nodes to construct the tree
38
39    // Calculate the middle index and use it as the root to balance the tree
40    const mid = start + Math.floor((end - start) / 2);
41    const newNode = new TreeNode(nodeValues[mid]); // Create a new node with the mid value
42
43    // Recursively construct the left and right subtrees and link them to the new node
44    newNode.left = rebuildTree(start, mid - 1);
45    newNode.right = rebuildTree(mid + 1, end);
46  
47    return newNode; // Return the new subtree root
48}
49

Time and Space Complexity

The given code consists of a depth-first search (DFS) traversal (dfs function) to collect the values of the nodes in an inorder fashion and a recursive function (build function) to construct a balanced binary search tree (BST) from the sorted list of values.

Time Complexity

  1. DFS Traversal (dfs function): The DFS traversal visits each node exactly once. With n being the number of nodes in the tree, the time complexity for the DFS traversal is O(n).

  2. Constructing Balanced BST (build function): The build function is called once for each element in the sorted list of values. Similar to binary search, at each recursive call, it splits the list into two halves until the base case is reached (when i > j). Therefore, constructing the balanced BST would also take O(n) time since each node is visited once during the construction.

Adding both complexities together, the overall time complexity of the algorithm is O(n).

Space Complexity

  1. DFS Traversal: The auxiliary space is used to store the inorder traversal of the BST in the vals list, which will contain n elements, thus requiring O(n) space.

  2. Constructing Balanced BST: The recursive build function will use additional space on the call stack. In the worst case, there will be O(log n) recursive calls on the stack at the same time, given that the newly constructed BST is balanced (hence the maximum height of the call stack corresponds to the height of a balanced BST).

Overall, considering both the space used by the list vals and the recursion call stack, the space complexity of the algorithm is O(n) for the list and O(log n) for the call stack, leading to O(n) space complexity when combined.

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


Fast Track Your Learning with Our Quick Skills Quiz:

Suppose k is a very large integer(2^64). Which of the following is the largest as n grows to infinity?


Recommended Readings


Got a question? Ask the Monster Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.

Coding Interview Strategies

Dive into our free, detailed pattern charts and company guides to understand what each company focuses on.

See Patterns

🪄