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;
}

Comment