Lintcode 87. Remove Node in Binary Search Tree

Description:

Given a root of Binary Search Tree with unique value for each node. Remove the node with given value. If there is no such a node with given value in the binary search tree, do nothing. You should keep the tree still a binary search tree after removal.
Example:
Given binary search tree:

Remove 3, you can either return:

Thinking:

The main idea is the same as insert a node in BST. The difference is that we should know how to re-construct the tree if root is removed. There are two ways as the question shows us: we can let root.right be the root and insert the left tree at the end of the right tree. Also, we can let root.left be the root. And insert the right tree at the end of the left tree.

Steps:

Check whether input is valid.
Decide use recursion or iteration? The same as insert a node in BST. It's better to use recursion.
Do we need to create a helper function? Yes, during the process of re-constructing the tree, we need a function to find the leftmost leafnode of the right tree.
How to do the recursion? Comparing the root.val and node.val. And call function(root.left/right, node).
Determine the return condition. Typically,root == null. And then we need to consider root >=< value, In root.val == value we should thinkroot.left == null and/or root.right == null. Since we need to re-construct the tree.

Code:

/**
 * Definition of TreeNode:
 * public class TreeNode {
 *     public int val;
 *     public TreeNode left, right;
 *     public TreeNode(int val) {
 *         this.val = val;
 *         this.left = this.right = null;
 *     }
 * }
 */
public class Solution {
    /**
     * @param root: The root of the binary search tree.
     * @param value: Remove the node with given value.
     * @return: The root of the binary search tree after removal.
     */
    public TreeNode removeNode(TreeNode root, int value) {
        // write your code here
        if(root == null) return root;
        if(root.val == value){
            if(root.left == null && root.right == null) return null;
            if(root.left == null) return root.right;
            if(root.right == null) return root.left;
            TreeNode node = findHelper(root.right);
            node.left = root.left;
            return root.right;
        }
        else if(root.val < value ){
           root.right =  removeNode(root.right,value);
        }
        else{
           root.left = removeNode(root.left,value);
        }
        return root;
    }
    private TreeNode findHelper(TreeNode node){
        if(node == null) return null;
        TreeNode current = node;
        if(current.left != null){
            current = current.left;
        }
        return current;
    }
}

Conclusion:

We do the recursion. We should know the what the function does, and the value the function returns. In thiis problem we should be careful about whether we should return remove(root.right, value) or root.right = remove(root.right,value)

Lintcode 87. Remove Node in Binary Search Tree