Problem description
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).
Solution:
- First step is to find the key in the tree. If it is not found. We are done. This step takes O(log(n))
- If the node has no right child, we can just set the child of its parent to be this node’s left child. (The leftness and rightness depends of parent node depends on whether this node is left or right child of its parent). This step takes O(1)
- If the node has right child, we can employ continous tree rotation operations to make its right child to be None. In details, after each rotation, if right child is not None, then right child’s left child become key node’s right child.
Code is shown below
# 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: TreeNode, key: int) -> TreeNode:
if root == None:
return root
parent = None
current = root
while current and current.val != key:
parent = current
if current.val > key:
current = current.left
else:
current = current.right
# key not found
if current == None:
return root
# left rotate tree using key node as pivot, untils its right child is None
while current.right:
temp = current.right.left
current.right.left = current
temp2 = current.right
current.right = temp
if parent == None:
root = temp2 #if key is root, reset the root
else:
if current == parent.left:
parent.left = temp2
else:
parent.right = temp2
parent = temp2 #reset parent node
if parent == None: #no rotation performed and key is root
return current.left
else: #after all rotation, key's right child is None
if current == parent.left:
parent.left = current.left
else:
parent.right = current.left
return root