LeetCode 450. Delete Node in a BST [Swift 题解]

题目

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).

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

附: 原题目链接

思路

第一步, 查找要删除的数值:
首先BST的特性, left比root数值小, right比root数值大, 所以先recursive的找到该数值,如果没有则返回原root

第二步, 删除找到的Node:
这是这个题目的关键地方,如何删除找到的Node, 重新construct 成合法的BST呢。 我们可以找到右边子树里最小的那个数(应该是在右边的子树里最底层左边的那个数字), 然后直接把那个Node删除掉,替换到我们开始要删除的Node即可。

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     public var val: Int
 *     public var left: TreeNode?
 *     public var right: TreeNode?
 *     public init(_ val: Int) {
 *         self.val = val
 *         self.left = nil
 *         self.right = nil
 *     }
 * }
 */
class Solution {  
    func deleteNode(_ root: TreeNode?, _ key: Int) -> TreeNode? {
        guard let root = root else {
            return nil
        }

        if root.val > key {
            root.left = deleteNode(root.left, key)
        } else if root.val < key {
            root.right = deleteNode(root.right, key)
        } else {
            if root.left == nil {
                return root.right
            } else if root.right == nil {
                return root.left
            } else {
                let temp = root

                let newRoot = findMin(root.right!)

                newRoot.right = deleteNode(root.right, newRoot.val)
                newRoot.left = temp.left
                return newRoot
            }
        }

        return root
    }

    func findMin(_ root: TreeNode) -> TreeNode {
        var _root = root
        while _root.left != nil {
            _root = _root.left!
        }

        return _root
    }
}

chenbo

编程爱好者

Singapore