Taro Logo

Balance a Binary Search Tree

Medium
2 months ago

Given the root of a binary search tree, return a balanced binary search tree with the same node values. A binary search tree is balanced if the depth of the two subtrees of every node never differs by more than 1. If there is more than one answer, return any of them. Let's illustrate with an example:

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.

Constraints:

  • The number of nodes in the tree is in the range [1, 10^4].
  • 1 <= Node.val <= 10^5

Write code to solve this problem.

Sample Answer
## Balanced Binary Search Tree

### Problem Description

Given the root of a binary search tree, return a balanced binary search tree with the same node values. A binary search tree is balanced if the depth of the two subtrees of every node never differs by more than 1. If there is more than one answer, return any of them.

**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]


### Naive Approach

The naive approach would involve collecting all the node values from the BST into an array, sorting the array, and then constructing a new balanced BST from the sorted array. This approach ensures that the resulting tree contains all the original values and is balanced since we are inserting the middle element as the root and recursively building the left and right subtrees.

**Code (Java):**

```java
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

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> values = new ArrayList<>();
        inorderTraversal(root, values);
        Collections.sort(values);
        return buildBalancedBST(values, 0, values.size() - 1);
    }

    private void inorderTraversal(TreeNode node, List<Integer> values) {
        if (node == null) {
            return;
        }
        inorderTraversal(node.left, values);
        values.add(node.val);
        inorderTraversal(node.right, values);
    }

    private TreeNode buildBalancedBST(List<Integer> values, int start, int end) {
        if (start > end) {
            return null;
        }
        int mid = (start + end) / 2;
        TreeNode node = new TreeNode(values.get(mid));
        node.left = buildBalancedBST(values, start, mid - 1);
        node.right = buildBalancedBST(values, mid + 1, end);
        return node;
    }
}

Optimal Approach

The optimal approach is very similar, but it skips the explicit sorting step by leveraging the fact that an inorder traversal of a BST already produces a sorted list of values. So, we perform an inorder traversal, store the values in a list, and then use that sorted list to build a balanced BST. This saves us the O(n log n) time complexity of sorting.

Code (Java):

import java.util.ArrayList;
import java.util.List;

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> values = new ArrayList<>();
        inorderTraversal(root, values);
        return buildBalancedBST(values, 0, values.size() - 1);
    }

    private void inorderTraversal(TreeNode node, List<Integer> values) {
        if (node == null) {
            return;
        }
        inorderTraversal(node.left, values);
        values.add(node.val);
        inorderTraversal(node.right, values);
    }

    private TreeNode buildBalancedBST(List<Integer> values, int start, int end) {
        if (start > end) {
            return null;
        }
        int mid = (start + end) / 2;
        TreeNode node = new TreeNode(values.get(mid));
        node.left = buildBalancedBST(values, start, mid - 1);
        node.right = buildBalancedBST(values, mid + 1, end);
        return node;
    }
}

Big O Runtime Analysis

  1. Inorder Traversal: O(n), where n is the number of nodes in the BST, because we visit each node once.
  2. Sorting (Naive): O(n log n), where n is the number of nodes, since we use Collections.sort().
  3. Build Balanced BST: O(n), where n is the number of nodes, because we visit each element in the sorted list once.

Therefore:

  • Naive Approach: O(n) + O(n log n) + O(n) = O(n log n)
  • Optimal Approach: O(n) + O(n) = O(n)

Big O Space Usage Analysis

  1. Values List: O(n), where n is the number of nodes, as we store all node values in a list.
  2. Recursion Stack: O(log n) in the average case for a balanced tree, and O(n) in the worst case for a skewed tree, due to recursive calls in buildBalancedBST.

Therefore, the overall space complexity is O(n) + O(log n) which simplifies to O(n), with O(n) dominating.

Edge Cases

  1. Empty Tree: If the input tree is empty (root is null), the algorithm should return null without any errors.
  2. Single Node Tree: If the tree contains only one node, the algorithm should return the same node, as it is already balanced.
  3. Skewed Tree: The algorithm should handle skewed trees efficiently by rebalancing them. For example, a tree that is heavily left-skewed or right-skewed.
  4. Duplicate Values: The algorithm should correctly handle trees with duplicate values. The inorder traversal will maintain the order of duplicates, and the balancing process will place them accordingly.
class Solution {
    public TreeNode balanceBST(TreeNode root) {
        if (root == null) {
            return null; // Handles empty tree
        }
        
        List<Integer> values = new ArrayList<>();
        inorderTraversal(root, values);
        return buildBalancedBST(values, 0, values.size() - 1);
    }

    private void inorderTraversal(TreeNode node, List<Integer> values) {
        if (node == null) {
            return;
        }
        inorderTraversal(node.left, values);
        values.add(node.val);
        inorderTraversal(node.right, values);
    }

    private TreeNode buildBalancedBST(List<Integer> values, int start, int end) {
        if (start > end) {
            return null;
        }
        int mid = (start + end) / 2;
        TreeNode node = new TreeNode(values.get(mid));
        node.left = buildBalancedBST(values, start, mid - 1);
        node.right = buildBalancedBST(values, mid + 1, end);
        return node;
    }
}

Explanation of Edge Case Handling in Code:

  • Empty Tree: The if (root == null) check at the beginning of the balanceBST function handles the case where the input tree is empty, returning null immediately.

  • Single Node Tree: The algorithm implicitly handles this as well. If the tree only has one node, then the inorder traversal results in a list of size one, and when buildBalancedBST is called, start and end both point to the same index. The mid is same as start/end, and left/right subtrees are null as the function recurses with start > end.

  • Skewed and Duplicate Values: The inorderTraversal always generates a sorted list of values which the buildBalancedBST function uses to create a balanced tree, regardless of how skewed the original tree was, or whether it had any duplicate values.