Delete Node in a BST

题目
Given a root node reference of a BST and a key, delete the node with the given key in the BST. Return the root node reference (possibly updated) of the BST.

Basically, the deletion can be divided into two stages:

Search for a node to remove.
If the node is found, delete the node.
Note: Time complexity should be O(height of tree).

答案
看起来有点长，或许可以优化一下

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {
    public TreeNode deleteNode(TreeNode root, int key) {
        TreeNode curr = root, parent = null, del_node = null, replace;
        boolean left = true;
        while(curr != null) {
            if(key == curr.val) {
                del_node = curr;
                break;
            }
            parent = curr;
            left = key < curr.val;
            curr = (key < curr.val) ? curr.left : curr.right;            
        }
        // If the node with val key is not found, just return the tree itself
        if(del_node == null) return root;

        replace = findReplaceNode(del_node.left);
        if(replace == null) {
            replace = del_node.right;
        }
        else {
            // Inherit del_node's children
            replace.right = del_node.right;
            if(replace != del_node.left) {
                replace.left = del_node.left;
            }            
        }
        if(left && parent != null) parent.left = replace;
        if(!left && parent != null) parent.right = replace;
        return (del_node == root)? replace : root;
    }

    private TreeNode findReplaceNode(TreeNode root) {
        if(root == null) return null;
        // Go right until it's not possible
        TreeNode curr = root, parent = null;
        while(curr.right != null) {
            parent = curr;
            curr = curr.right;
        }

        if(curr == root) return curr;
        parent.right = curr.left;
        return curr;
    }
}

    public TreeNode deleteNode(TreeNode root, int key) {
        if(root == null) return null;
        if(root.val == key){
            if(root.right == null){
                return root.left;
            }else{
                TreeNode ptr = root.right;
                while(ptr.left!=null){
                    ptr = ptr.left;
                }
                ptr.left = root.left;
                return root.right;
            }
        }else if(root.val < key){
            root.right = deleteNode(root.right, key);
        }else if(root.val > key){
            root.left = deleteNode(root.left, key);
        }
        return root;
    }