450. Delete Node in a BST - Medium

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:

1. Search for a node to remove.
2. If the node is found, delete the node.

 

Note: Time complexity should be O(height of tree).

Example:

root = [5,3,6,2,4,null,7]
key = 3

    5
   / \
  3   6
 / \   \
2   4   7

Given key to delete is 3. So we find the node with value 3 and delete it.

One valid answer is [5,4,6,2,null,null,7], shown in the following BST.

    5
   / \
  4   6
 /     \
2       7

Another valid answer is [5,2,6,null,4,null,7].

    5
   / \
  2   6
   \   \
    4   7

 

1. if root.val == key, need to delete root
　　1.1 if root only has right child, just return root.right
　　1.2 if root only has left child, just return root.left
　　1.3 if root has two child nodes, need to find a candidate node in right subtree and make it new root, the candidate is usually root's right child, and also need to modify the tree, make root's left subtree to be candidate's left subtree
　　　　1.3.1 if root.right.left == null, just move root's left subtree to root.right's left, and return root.right as new root
　　　　1.3.2 if root.right.left is not null, need to find the smallest node in root's right subtree as candidate new root, delete it from its original place, move root's left/right subtree to smallest's left/right subtree, and then return smallest as new root

2. if root.val > key, use recursion
3. if root.val < key, use recursion

helper function: deleteSmallest(node)
go all the way left to node until node.left.left == null, smallest node is node.left,
delete smallest node by moving smallest node's right subtree to node's left, and return the smallest node

time = O(height), space = O(height)

class Solution {
    public TreeNode deleteNode(TreeNode root, int key) {
        if(root == null) {
            return root;
        }
        if(root.val == key) {
            if(root.left == null) {
                return root.right;
            } else if(root.right == null) {
                return root.left;
            } else if(root.right.left == null) {
                root.right.left = root.left;
                return root.right;
            } else {
                TreeNode smallest = deleteSmallest(root.right);
                smallest.left = root.left;
                smallest.right = root.right;
                return smallest;
            }
        }
        
        if(root.val > key) {
            root.left = deleteNode(root.left, key);
        } else if(root.val < key) {
            root.right = deleteNode(root.right, key);
        }
        
        return root;
    }
    
    private TreeNode deleteSmallest(TreeNode node) {
        while(node.left.left != null) {
            node = node.left;
        }
        TreeNode smallest = node.left;
        node.left = node.left.right;
        return smallest;
    }
}