Trim a Binary Search Tree

Given the root of a binary search tree and the lowest and highest boundaries as low and high, trim the tree so that all its elements lies in [low, high]. Trimming the tree should not change the relative structure of the elements that will remain in the tree (i.e., any node's descendant should remain a descendant). It can be proven that there is a unique answer.

Return the root of the trimmed binary search tree. Note that the root may change depending on the given bounds.

Example 1: Input: root = [1,0,2], low = 1, high = 2 Output: [1,null,2] Example 2: Input: root = [3,0,4,null,2,null,null,1], low = 1, high = 3 Output: [3,2,null,1] Example 3: Input: root = [1], low = 1, high = 2 Output: [1] Example 4: Input: root = [1,null,2], low = 1, high = 3 Output: [1,null,2] Example 5: Input: root = [1,null,2], low = 2, high = 4 Output: [2] Constraints:

• The number of nodes in the tree in the range [1, 104].

• 0 <= Node.val <= 104

• The value of each node in the tree is unique.

• root is guaranteed to be a valid binary search tree.

• 0 <= low <= high <= 104

Solution:

Iterative:

public TreeNode trimBST(TreeNode root, int low, int high) {
    if (root == null) {
      return null;
    }

    TreeNode dummyHead = new TreeNode();
    Stack<List<TreeNode>> stack = new Stack<>();
    stack.push(Arrays.asList(root, dummyHead));
    TreeNode currentNode;
    List<TreeNode> pair;
    while (!stack.isEmpty()) {
      pair = stack.pop();
      currentNode = pair.get(0);
      if (currentNode.val >= high) {
        currentNode.right = null;
      }
      if (currentNode.val <= low) {
        currentNode.left = null;
      }
      if (currentNode.val > low && currentNode.left != null) {
        stack.push(Arrays.asList(currentNode.left, currentNode.val > high ? pair.get(1) : currentNode));
      }
      if (currentNode.val < high && currentNode.right != null) {
        stack.push(Arrays.asList(currentNode.right, currentNode.val < low ? pair.get(1) : currentNode));
      }
      if (pair.get(1) == dummyHead || currentNode.val < pair.get(1).val) {
        pair.get(1).left = currentNode;
      } else {
        pair.get(1).right = currentNode;
      }
    }
    return dummyHead.left;
  }

Recursive:

class Solution {
    public TreeNode trimBST(TreeNode root, int low, int high) {
         if (root == null) return root;
        if (root.val > high) return trimBST(root.left, low, high);
        if (root.val < low) return trimBST(root.right, low, high);

        root.left = trimBST(root.left, low, high);
        root.right = trimBST(root.right, low, high);
        return root;   
    }
}