450. Delete Node in a BST

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:

还有bug

class Solution(object):
    def deleteNode(self, root, key):
        """
        :type root: TreeNode
        :type key: int
        :rtype: TreeNode
        """
        if not root: return root
        if key<root.val: 
            root.left = self.deleteNode(root.left, key)
            return root
        elif key>root.val: 
            root.right = self.deleteNode(root.right, key)
            return root
        else:
            if not root.left: return root.right
            elif not root.left.right: 
                root.left.right = root.right
                return root.left
            else:
                pre, lm = root.left, root.left.right
                while lm and lm.right: pre, lm=lm, lm.right
                pre.right = None
                lm.left = root.left
                lm.right = root.right
                return lm