776. Split BST

Given the root of a binary search tree (BST) and an integer target, split the tree into two subtrees where one subtree has nodes that are all smaller or equal to the target value, while the other subtree has all nodes that are greater than the target value. It Is not necessarily the case that the tree contains a node with the value target.

Additionally, most of the structure of the original tree should remain. Formally, for any child c with parent p in the original tree, if they are both in the same subtree after the split, then node c should still have the parent p.

Return an array of the two roots of the two subtrees.

Example 1:

Input: root = [4,2,6,1,3,5,7], target = 2
Output: [[2,1],[4,3,6,null,null,5,7]]

Example 2:

Input: root = [1], target = 1
Output: [[1],[]]

Constraints:

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

• 0 <= Node.val, target <= 1000

Solution

We are going to split the tree into two trees. One is greaterTree which all nodes in the tree are greater than the target. Another is smallerTree which all nodes are smaller or equal to the target.

We noticed that if node.val > target, then node and node.right must be in the greaterTree. But for it left child node, we need a furter check by splitBST(node.left), we will get greaterTree2 and smallerTree2 by splitBST(node.left). The greaterTree2 will be the left child of the greaterTree. And the smallerTree2 will be the smallerTree of splitBST(node).

And if the node.val <= target, the node and node.left shall be in the smaller tree, and we also need to check node.left by splitBST(node.right). we will get greaterTree2 and smallerTree2 by splitBST(node.right). The smallerTree2 will be the smallerTree.right. And The greaterTree2 is the greaterTree of splitBST(node)

/**
 * Definition for a binary tree node.
 * public 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[] splitBST(TreeNode root, int target) {
        TreeNode greater = null;
        TreeNode smaller = null;
        if(root != null && root.val <= target){
            smaller = root;
            TreeNode[] right = splitBST(root.right, target);
            smaller.right = right[0];
            greater = right[1];
            //root and it's left child shall be splited
            //for it's right child, we need extra check
        } else if(root != null && root.val > target){
            greater = root;
            TreeNode[] left = splitBST(root.left, target);
            greater.left = left[1];
            smaller = left[0];
        }
        return new TreeNode[]{smaller, greater};
    }
}