LeetCode in Kotlin

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.

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: []

Constraints:

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

Solution

import com_github_leetcode.TreeNode

/*
 * Example:
 * var ti = TreeNode(5)
 * var v = ti.`val`
 * Definition for a binary tree node.
 * class TreeNode(var `val`: Int) {
 *     var left: TreeNode? = null
 *     var right: TreeNode? = null
 * }
 */
class Solution {
    fun deleteNode(root: TreeNode?, key: Int): TreeNode? {
        if (root == null) return root
        // find the correct node
        if (key < root.`val`) {
            root.left = deleteNode(root.left, key)
        } else if (key > root.`val`) {
            root.right = deleteNode(root.right, key)
        } else {
            // case 1 - both children are null
            if (root.left == null && root.right == null) {
                return null
            } else if (root.left != null && root.right != null) { // case 2 - both children are NOT null
                val successor = minimum(root.right!!) // inorder successor
                root.right = deleteNode(root.right, successor)
                root.`val` = successor
            } else { // case 3 - only one of the child is null
                return root.left ?: root.right
            }
        }
        return root
    }

    private fun minimum(root: TreeNode): Int {
        var node = root
        while (node.left != null) {
            node = node.left!!
        }
        return node.`val`
    }
}