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.
Follow up: Can you solve it with time complexity O(height of tree)?
- code
class Solution(object):
def deleteNode(self, root, key):
"""
:type root: TreeNode
:type key: int
:rtype: TreeNode
"""
if not root: return None
if root.val == key:
if root.left:
# Find the right most leaf of the left sub-tree
left_right_most = root.left
while left_right_most.right:
left_right_most = left_right_most.right
# Attach right child to the right of that leaf
left_right_most.right = root.right
# Return left child instead of root, a.k.a delete root
return root.left
else:
return root.right
# If left or right child got deleted, the returned root is the child of the deleted node.
elif root.val > key:
root.left = self.deleteNode(root.left, key)
else:
root.right = self.deleteNode(root.right, key)
return root