[leetcode] 450. Delete Node in a BST @ python

原题

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

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
Accepted
53,150
Submissions
136,957

解法

递归. Base case是如果root为空, 返回None. 然后使用分治法: 如果key比root值小, 那么从左边子树中删除key; 如果key比root值大, 那么从右边子树中删除key; 如果key等于root值, 再分为3种情况: 如果root没有左子树, 则返回右子树; 如果root没有右子树, 则返回左子树; 如果root既有左子树也有右子树, 那么我们在右子树中找到值最小的节点去替换root, 然后删除最小值的那个节点.
Time: O(h), h为二叉树的高度
Space: O(1)

代码

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def deleteNode(self, root, key):
        """
        :type root: TreeNode
        :type key: int
        :rtype: TreeNode
        """
        # base case
        if not root: return None
        # divide and conquer
        if root.val > key:
            root.left = self.deleteNode(root.left, key)
        elif root.val < key:
            root.right = self.deleteNode(root.right, key)
        else:
            if not root.left:
                return root.right
            if not root.right:
                return root.left
            else:
                # find the minimun node in the right subtree and replace the root with that node
                temp = root.right
                mini = temp.val
                while temp.left:
                    temp = temp.left
                    mini = temp.val
                # replace the value of root with minimun val
                root.val = mini
                root.right = self.deleteNode(root.right, root.val)
        return root