Improved del_node func

There was some logical error in implementation of delete node function.

Also we don't need to find balance factor 2 times so made separate variable.
This commit is contained in:
Shaurya Bisht 2024-10-12 19:54:55 +05:30
parent e9e7c96465
commit 0719a18d47

View File

@ -196,8 +196,13 @@ def get_left_most(root: MyNode) -> Any:
def del_node(root: MyNode, data: Any) -> MyNode | None:
if root is None:
print("Nothing to delete")
return None
left_child = root.get_left()
right_child = root.get_right()
if root.get_data() == data:
if left_child is not None and right_child is not None:
temp_data = get_left_most(right_child)
@ -221,20 +226,23 @@ def del_node(root: MyNode, data: Any) -> MyNode | None:
else:
root.set_right(del_node(right_child, data))
if get_height(right_child) - get_height(left_child) == 2:
root.set_height(my_max(get_height(root.get_right()), get_height(root.get_left())) + 1)
balance_factor = get_height(root.get_left()) - get_height(root.get_right())
if balance_factor == 2:
assert right_child is not None
if get_height(right_child.get_right()) > get_height(right_child.get_left()):
root = left_rotation(root)
else:
root = rl_rotation(root)
elif get_height(right_child) - get_height(left_child) == -2:
elif balance_factor == -2:
assert left_child is not None
if get_height(left_child.get_left()) > get_height(left_child.get_right()):
root = right_rotation(root)
else:
root = lr_rotation(root)
height = my_max(get_height(root.get_right()), get_height(root.get_left())) + 1
root.set_height(height)
return root