Skip to content

Commit 0719a18

Browse files
committed
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.
1 parent e9e7c96 commit 0719a18

File tree

1 file changed

+12
-4
lines changed

1 file changed

+12
-4
lines changed

data_structures/binary_tree/avl_tree.py

+12-4
Original file line numberDiff line numberDiff line change
@@ -196,8 +196,13 @@ def get_left_most(root: MyNode) -> Any:
196196

197197

198198
def del_node(root: MyNode, data: Any) -> MyNode | None:
199+
if root is None:
200+
print("Nothing to delete")
201+
return None
202+
199203
left_child = root.get_left()
200204
right_child = root.get_right()
205+
201206
if root.get_data() == data:
202207
if left_child is not None and right_child is not None:
203208
temp_data = get_left_most(right_child)
@@ -221,20 +226,23 @@ def del_node(root: MyNode, data: Any) -> MyNode | None:
221226
else:
222227
root.set_right(del_node(right_child, data))
223228

224-
if get_height(right_child) - get_height(left_child) == 2:
229+
root.set_height(my_max(get_height(root.get_right()), get_height(root.get_left())) + 1)
230+
231+
balance_factor = get_height(root.get_left()) - get_height(root.get_right())
232+
233+
if balance_factor == 2:
225234
assert right_child is not None
226235
if get_height(right_child.get_right()) > get_height(right_child.get_left()):
227236
root = left_rotation(root)
228237
else:
229238
root = rl_rotation(root)
230-
elif get_height(right_child) - get_height(left_child) == -2:
239+
elif balance_factor == -2:
231240
assert left_child is not None
232241
if get_height(left_child.get_left()) > get_height(left_child.get_right()):
233242
root = right_rotation(root)
234243
else:
235244
root = lr_rotation(root)
236-
height = my_max(get_height(root.get_right()), get_height(root.get_left())) + 1
237-
root.set_height(height)
245+
238246
return root
239247

240248

0 commit comments

Comments
 (0)