avl_tree.py

"""
Implementation of an auto-balanced binary tree!
For doctests run following command:
python3 -m doctest -v avl_tree.py
For testing run:
python avl_tree.py
"""

from __future__ importannotations

importmath
importrandom
fromtypingimportAny


classMyQueue:
def__init__(self) ->None:
self.data: list[Any] = []
self.head: int=0
self.tail: int=0

defis_empty(self) ->bool:
returnself.head==self.tail

defpush(self, data: Any) ->None:
self.data.append(data)
self.tail=self.tail+1

defpop(self) ->Any:
ret=self.data[self.head]
self.head=self.head+1
returnret

defcount(self) ->int:
returnself.tail-self.head

defprint_queue(self) ->None:
print(self.data)
print("**************")
print(self.data[self.head : self.tail])


classMyNode:
def__init__(self, data: Any) ->None:
self.data=data
self.left: MyNode|None=None
self.right: MyNode|None=None
self.height: int=1

defget_data(self) ->Any:
returnself.data

defget_left(self) ->MyNode|None:
returnself.left

defget_right(self) ->MyNode|None:
returnself.right

defget_height(self) ->int:
returnself.height

defset_data(self, data: Any) ->None:
self.data=data

defset_left(self, node: MyNode|None) ->None:
self.left=node

defset_right(self, node: MyNode|None) ->None:
self.right=node

defset_height(self, height: int) ->None:
self.height=height


defget_height(node: MyNode|None) ->int:
ifnodeisNone:
return0
returnnode.get_height()


defmy_max(a: int, b: int) ->int:
ifa>b:
returna
returnb


defright_rotation(node: MyNode) ->MyNode:
r"""
 A B
 / \ / \
 B C Bl A
 / \ --> / / \
 Bl Br UB Br C
 /
 UB
 UB = unbalanced node
 """
print("left rotation node:", node.get_data())
ret=node.get_left()
assertretisnotNone
node.set_left(ret.get_right())
ret.set_right(node)
h1=my_max(get_height(node.get_right()), get_height(node.get_left())) +1
node.set_height(h1)
h2=my_max(get_height(ret.get_right()), get_height(ret.get_left())) +1
ret.set_height(h2)
returnret


defleft_rotation(node: MyNode) ->MyNode:
"""
 a mirror symmetry rotation of the left_rotation
 """
print("right rotation node:", node.get_data())
ret=node.get_right()
assertretisnotNone
node.set_right(ret.get_left())
ret.set_left(node)
h1=my_max(get_height(node.get_right()), get_height(node.get_left())) +1
node.set_height(h1)
h2=my_max(get_height(ret.get_right()), get_height(ret.get_left())) +1
ret.set_height(h2)
returnret


deflr_rotation(node: MyNode) ->MyNode:
r"""
 A A Br
 / \ / \ / \
 B C LR Br C RR B A
 / \ --> / \ --> / / \
 Bl Br B UB Bl UB C
 \ /
 UB Bl
 RR = right_rotation LR = left_rotation
 """
left_child=node.get_left()
assertleft_childisnotNone
node.set_left(left_rotation(left_child))
returnright_rotation(node)


defrl_rotation(node: MyNode) ->MyNode:
right_child=node.get_right()
assertright_childisnotNone
node.set_right(right_rotation(right_child))
returnleft_rotation(node)


definsert_node(node: MyNode|None, data: Any) ->MyNode|None:
ifnodeisNone:
returnMyNode(data)
ifdata<node.get_data():
node.set_left(insert_node(node.get_left(), data))
if (
get_height(node.get_left()) -get_height(node.get_right()) ==2
 ): # an unbalance detected
left_child=node.get_left()
assertleft_childisnotNone
if (
data<left_child.get_data()
 ): # new node is the left child of the left child
node=right_rotation(node)
else:
node=lr_rotation(node)
else:
node.set_right(insert_node(node.get_right(), data))
ifget_height(node.get_right()) -get_height(node.get_left()) ==2:
right_child=node.get_right()
assertright_childisnotNone
ifdata<right_child.get_data():
node=rl_rotation(node)
else:
node=left_rotation(node)
h1=my_max(get_height(node.get_right()), get_height(node.get_left())) +1
node.set_height(h1)
returnnode


defget_right_most(root: MyNode) ->Any:
whileTrue:
right_child=root.get_right()
ifright_childisNone:
break
root=right_child
returnroot.get_data()


defget_left_most(root: MyNode) ->Any:
whileTrue:
left_child=root.get_left()
ifleft_childisNone:
break
root=left_child
returnroot.get_data()


defdel_node(root: MyNode, data: Any) ->MyNode|None:
left_child=root.get_left()
right_child=root.get_right()
ifroot.get_data() ==data:
ifleft_childisnotNoneandright_childisnotNone:
temp_data=get_left_most(right_child)
root.set_data(temp_data)
root.set_right(del_node(right_child, temp_data))
elifleft_childisnotNone:
root=left_child
elifright_childisnotNone:
root=right_child
else:
returnNone
elifroot.get_data() >data:
ifleft_childisNone:
print("No such data")
returnroot
else:
root.set_left(del_node(left_child, data))
# root.get_data() < data
elifright_childisNone:
returnroot
else:
root.set_right(del_node(right_child, data))

# Re-fetch left_child and right_child references
left_child=root.get_left()
right_child=root.get_right()

ifget_height(right_child) -get_height(left_child) ==2:
assertright_childisnotNone
ifget_height(right_child.get_right()) >get_height(right_child.get_left()):
root=left_rotation(root)
else:
root=rl_rotation(root)
elifget_height(right_child) -get_height(left_child) ==-2:
assertleft_childisnotNone
ifget_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)
returnroot


classAVLtree:
"""
 An AVL tree doctest
 Examples:
 >>> t = AVLtree()
 >>> t.insert(4)
 insert:4
 >>> print(str(t).replace(" \\n","\\n"))
 4
 *************************************
 >>> t.insert(2)
 insert:2
 >>> print(str(t).replace(" \\n","\\n").replace(" \\n","\\n"))
 4
 2 *
 *************************************
 >>> t.insert(3)
 insert:3
 right rotation node: 2
 left rotation node: 4
 >>> print(str(t).replace(" \\n","\\n").replace(" \\n","\\n"))
 3
 2 4
 *************************************
 >>> t.get_height()
 2
 >>> t.del_node(3)
 delete:3
 >>> print(str(t).replace(" \\n","\\n").replace(" \\n","\\n"))
 4
 2 *
 *************************************
 """

def__init__(self) ->None:
self.root: MyNode|None=None

defget_height(self) ->int:
returnget_height(self.root)

definsert(self, data: Any) ->None:
print("insert:"+str(data))
self.root=insert_node(self.root, data)

defdel_node(self, data: Any) ->None:
print("delete:"+str(data))
ifself.rootisNone:
print("Tree is empty!")
return
self.root=del_node(self.root, data)

def__str__(
self,
 ) ->str: # a level traversale, gives a more intuitive look on the tree
output=""
q=MyQueue()
q.push(self.root)
layer=self.get_height()
iflayer==0:
returnoutput
cnt=0
whilenotq.is_empty():
node=q.pop()
space=" "*int(math.pow(2, layer-1))
output+=space
ifnodeisNone:
output+="*"
q.push(None)
q.push(None)
else:
output+=str(node.get_data())
q.push(node.get_left())
q.push(node.get_right())
output+=space
cnt=cnt+1
foriinrange(100):
ifcnt==math.pow(2, i) -1:
layer=layer-1
iflayer==0:
output+="\n*************************************"
returnoutput
output+="\n"
break
output+="\n*************************************"
returnoutput


def_test() ->None:
importdoctest

doctest.testmod()


if__name__=="__main__":
_test()
t=AVLtree()
lst=list(range(10))
random.shuffle(lst)
foriinlst:
t.insert(i)
print(str(t))
random.shuffle(lst)
foriinlst:
t.del_node(i)
print(str(t))