Reference: LeetCode
Difficulty: Easy

Problem

Given a binary search tree and the lowest and highest boundaries as L and R, trim the tree so that all its elements lies in [L, R] (R >= L). You might need to change the root of the tree, so the result should return the new root of the trimmed binary search tree.

Note:

Example: 

1
2
3
4
5
6
7
8
9
10
11
12
Input: 
    1
   / \
  0   2

  L = 1
  R = 2

Output: 
    1
      \
       2
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
Input: 
    3
   / \
  0   4
   \
    2
   /
  1

  L = 1
  R = 3

Output: 
      3
     / 
   2   
  /
 1

Follow up: Provide an iterative solution that take $O(1)$ space.

Analysis

Methods:

I tried using inorder traversal, but it does not work well because of the link connection.

  1. Recursion
    • If the root is not within [L, R], we should replace it by the left or right child. (2 cases)
    • If the root is within the range, we should keep the root node.
    • Time: $O(N)$ since each node is traversed.
    • Space: $O(h)$
  2. Iteration
    • Iteration solution is tricky.
    • We first find the root that is within the range and replace the invalid one.
    • Then we check the children and replace them if necessary.
    • Time: $O(N)$
    • Space: $O(1)$

Code

Test Case:

1
2
3
[1,0,2]
3
4

Recursion

Note:

  • else is necessary since we have updated the root. We have to check the root again in the next trimBST call.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
public TreeNode trimBST(TreeNode root, int L, int R) {
  if (root == null) {
    return null;
  }
  // trim
  if (root.val < L || root.val > R) {
    // trim root node
    root = (root.val < L) ? root.right : root.left; // cut
    root = trimBST(root, L, R);
  } else {
    // root in [L, R]
    root.left = trimBST(root.left, L, R);
    root.right = trimBST(root.right, L, R);
  }
  return root;
}

Another version:

1
2
3
4
5
6
7
8
9
10
11
12
public TreeNode trimBST(TreeNode root, int L, int R) {
  if (root == null) {
    return null;
  }
  // val not in range, trim
  if (root.val < L) return trimBST(root.right, L, R);
  if (root.val > R) return trimBST(root.left, L, R);
  // val in [L, R] (done with root)
  root.left = trimBST(root.left, L, R);
  root.right = trimBST(root.right, L, R);
  return root;
}

Iteration

Reference: Java solution, iteration version

The code in the post has a null bug in the code of finding the root. Try out this test case to see the error.

1
2
3
[1,0,2]
3
4

Note:

  • When looking for the new root, we may go left or right, so we can’t just separate two directions in two while loops.
  • Draw a picture and you can get a more intuitive picture why we handle the children like this. Hone in the properties of BST.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
public TreeNode trimBST(TreeNode root, int L, int R) {
  // find the okay root
  while (root != null && (root.val < L || root.val > R)) {
    if (root.val < L) {
      root = root.right;
    } else if (root.val > R) {
      root = root.left;
    }
  } // now the root is finalized.
  
  // find the left child
  TreeNode curr = root;
  while (curr != null) {
    while (curr.left != null && curr.left.val < L) { // if curr.left.val > L: STOP
      curr.left = curr.left.right;
    }
    curr = curr.left;
  }
  
  // find the right child
  curr = root;
  while (curr != null) {
    while (curr.right != null && curr.right.val > R) {
      curr.right = curr.right.left;
    }
    curr = curr.right;
  }
  return root;
}