450. Delete Node in a BST

LeetCode 450. Delete Node in a BST

Description

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.

Follow up: Can you solve it with time complexity O(height of tree)?

Example 1:

Input: root = [5,3,6,2,4,null,7], key = 3
Output: [5,4,6,2,null,null,7]
Explanation: 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 above BST.
Please notice that another valid answer is [5,2,6,null,4,null,7] and it's also accepted.

Constraints:

The number of nodes in the tree is in the range [0, 104].
-105 <= Node.val <= 105
Each node has a unique value.
root is a valid binary search tree.
-105 <= key <= 105

Solution

Return null when root is null. We find the target node first, then check if it matches the following scenario:

The target is a leaf, then return null;
The target only has one child, then return that child node;
The target has both left and right children. We first find the smallest (leftmost) node of the right sub-tree overwrite the key (root's value) with successor's value. The next step is to delete the new key - successor's value from the sub-tree whose root is the successor, and assign the processed tree root to the ancestor's child.

Complexity

Time complexity: $O(\log(n))$ , $\log(n)=H$ when the tree is a balance one;
Space complexity: $O(H)$

Code

func deleteNode(root *TreeNode, key int) *TreeNode {
	if root == nil {
		return root
	}
	if key > root.Val {
		root.Right = deleteNode(root.Right, key)
	} else if key < root.Val {
		root.Left = deleteNode(root.Left, key)
	} else { // found target
		// 1. leaf: delete itself
		if root.Left == nil && root.Right == nil {
			return nil
		}
		// 2. only left is not null, keep left
		if root.Left != nil && root.Right == nil {
			return root.Left
		}
		// 3. only right is not null, keep right
		if root.Left == nil && root.Right != nil {
			return root.Right
		}
		// 4. both are not null: find the smallest node of the right subtree
		ancestor, successor := root, root.Right
		for successor.Left != nil {
			ancestor = successor
			successor = successor.Left
		}
		root.Val = successor.Val        // overwrite the key with successor
		if successor == ancestor.Left { // delete successor (duplicate)
			ancestor.Left = deleteNode(successor, successor.Val)
		} else {
			ancestor.Right = deleteNode(successor, successor.Val)
		}
	}
	return root
}