Delete Node in a BST

Problem Statement

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.

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.

Example 2:

Input: root = [5,3,6,2,4,null,7], key = 0
Output: [5,3,6,2,4,null,7]
Explanation: The tree does not contain a node with value = 0.

Example 3:

Input: root = [], key = 0
Output: []

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

Solution

Idea: If we find the key node, we can move the entire right subtree to the end of the left subtree, or we can move the entire left subtree into the right subtree. We attach the subtrees at the end, which is the leaf node, and it still maintains the binary search tree property.

If the root is the key node, we can return a new node by attaching either the left or right subtree at the leaf node.

Time - O(n)

Space - O(1)

class Solution {
    private TreeNode helper(TreeNode node){
        if(node.left==null) return node.right;
        if(node.right==null) return node.left;
        TreeNode leftSubtree = node.left;
        TreeNode rightSubtree = node.right;
        TreeNode rightSubtreeLeftLeafNode = leftLeafNode(rightSubtree);
        rightSubtreeLeftLeafNode.left = leftSubtree;
        return rightSubtree;
    }
    private TreeNode leftLeafNode(TreeNode node){
        if(node == null || node.left==null){
            return node;
        }
        return leftLeafNode(node.left);
    }
    public TreeNode deleteNode(TreeNode root, int key) {

        if(root==null){
            return root;
        }

        if(root.val == key){
            return helper(root);
        }

        TreeNode temp = root;
        while(temp!=null){
            if(temp.val < key){
                if(temp.right!=null && temp.right.val == key){
                    temp.right = helper(temp.right);
                    return root;
                }
                temp = temp.right;
            }
            else if(temp.val > key){
                if(temp.left!=null && temp.left.val == key){
                    temp.left = helper(temp.left);
                    return root;
                }
                temp = temp.left;
            }
        }
        return root;
    }
}