Python/data_structures/binary_tree/avl_tree.py

349 lines
9.4 KiB
Python
Raw Normal View History

2019-10-05 05:14:13 +00:00
"""
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
2019-10-05 05:14:13 +00:00
"""
from __future__ import annotations
2018-11-19 19:34:44 +00:00
import math
import random
from typing import Any
2019-10-05 05:14:13 +00:00
class MyQueue:
def __init__(self) -> None:
self.data: list[Any] = []
self.head: int = 0
self.tail: int = 0
2019-10-05 05:14:13 +00:00
def is_empty(self) -> bool:
2018-11-19 19:34:44 +00:00
return self.head == self.tail
2019-10-05 05:14:13 +00:00
def push(self, data: Any) -> None:
2018-11-19 19:34:44 +00:00
self.data.append(data)
self.tail = self.tail + 1
2019-10-05 05:14:13 +00:00
def pop(self) -> Any:
2018-11-19 19:34:44 +00:00
ret = self.data[self.head]
self.head = self.head + 1
return ret
2019-10-05 05:14:13 +00:00
def count(self) -> int:
2018-11-19 19:34:44 +00:00
return self.tail - self.head
2019-10-05 05:14:13 +00:00
def print_queue(self) -> None:
2018-11-19 19:34:44 +00:00
print(self.data)
print("**************")
2019-10-05 05:14:13 +00:00
print(self.data[self.head : self.tail])
class MyNode:
def __init__(self, data: Any) -> None:
2018-11-19 19:34:44 +00:00
self.data = data
self.left: MyNode | None = None
self.right: MyNode | None = None
self.height: int = 1
2019-10-05 05:14:13 +00:00
def get_data(self) -> Any:
2018-11-19 19:34:44 +00:00
return self.data
2019-10-05 05:14:13 +00:00
def get_left(self) -> MyNode | None:
2018-11-19 19:34:44 +00:00
return self.left
2019-10-05 05:14:13 +00:00
def get_right(self) -> MyNode | None:
2018-11-19 19:34:44 +00:00
return self.right
2019-10-05 05:14:13 +00:00
def get_height(self) -> int:
2018-11-19 19:34:44 +00:00
return self.height
2019-10-05 05:14:13 +00:00
def set_data(self, data: Any) -> None:
2018-11-19 19:34:44 +00:00
self.data = data
return
2019-10-05 05:14:13 +00:00
def set_left(self, node: MyNode | None) -> None:
2018-11-19 19:34:44 +00:00
self.left = node
return
2019-10-05 05:14:13 +00:00
def set_right(self, node: MyNode | None) -> None:
2018-11-19 19:34:44 +00:00
self.right = node
return
2019-10-05 05:14:13 +00:00
def set_height(self, height: int) -> None:
2018-11-19 19:34:44 +00:00
self.height = height
return
2019-10-05 05:14:13 +00:00
def get_height(node: MyNode | None) -> int:
2018-11-19 19:34:44 +00:00
if node is None:
return 0
return node.get_height()
2018-11-19 19:34:44 +00:00
2019-10-05 05:14:13 +00:00
def my_max(a: int, b: int) -> int:
2018-11-19 19:34:44 +00:00
if a > b:
return a
return b
def right_rotation(node: MyNode) -> MyNode:
2019-10-05 05:14:13 +00:00
r"""
2018-11-22 06:33:50 +00:00
A B
/ \ / \
B C Bl A
/ \ --> / / \
Bl Br UB Br C
/
UB
UB = unbalanced node
2019-10-05 05:14:13 +00:00
"""
print("left rotation node:", node.get_data())
ret = node.get_left()
assert ret is not None
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)
2018-11-19 19:34:44 +00:00
return ret
2019-10-05 05:14:13 +00:00
def left_rotation(node: MyNode) -> MyNode:
2019-10-05 05:14:13 +00:00
"""
a mirror symmetry rotation of the left_rotation
2019-10-05 05:14:13 +00:00
"""
print("right rotation node:", node.get_data())
ret = node.get_right()
assert ret is not None
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)
2018-11-19 19:34:44 +00:00
return ret
2019-10-05 05:14:13 +00:00
def lr_rotation(node: MyNode) -> MyNode:
2019-10-05 05:14:13 +00:00
r"""
A A Br
2018-11-22 06:33:50 +00:00
/ \ / \ / \
B C LR Br C RR B A
2018-11-22 06:33:50 +00:00
/ \ --> / \ --> / / \
Bl Br B UB Bl UB C
2018-11-22 06:33:50 +00:00
\ /
UB Bl
RR = right_rotation LR = left_rotation
2019-10-05 05:14:13 +00:00
"""
left_child = node.get_left()
assert left_child is not None
node.set_left(left_rotation(left_child))
return right_rotation(node)
2018-11-22 06:33:50 +00:00
2019-10-05 05:14:13 +00:00
def rl_rotation(node: MyNode) -> MyNode:
right_child = node.get_right()
assert right_child is not None
node.set_right(right_rotation(right_child))
return left_rotation(node)
2018-11-19 19:34:44 +00:00
def insert_node(node: MyNode | None, data: Any) -> MyNode | None:
2018-11-19 19:34:44 +00:00
if node is None:
return MyNode(data)
if data < node.get_data():
node.set_left(insert_node(node.get_left(), data))
2019-10-05 05:14:13 +00:00
if (
get_height(node.get_left()) - get_height(node.get_right()) == 2
2019-10-05 05:14:13 +00:00
): # an unbalance detected
left_child = node.get_left()
assert left_child is not None
2019-10-05 05:14:13 +00:00
if (
data < left_child.get_data()
2019-10-05 05:14:13 +00:00
): # new node is the left child of the left child
node = right_rotation(node)
2018-11-19 19:34:44 +00:00
else:
node = lr_rotation(node)
2018-11-19 19:34:44 +00:00
else:
node.set_right(insert_node(node.get_right(), data))
if get_height(node.get_right()) - get_height(node.get_left()) == 2:
right_child = node.get_right()
assert right_child is not None
if data < right_child.get_data():
node = rl_rotation(node)
2018-11-19 19:34:44 +00:00
else:
node = left_rotation(node)
h1 = my_max(get_height(node.get_right()), get_height(node.get_left())) + 1
node.set_height(h1)
2018-11-19 19:34:44 +00:00
return node
2019-10-05 05:14:13 +00:00
def get_right_most(root: MyNode) -> Any:
while True:
right_child = root.get_right()
if right_child is None:
break
root = right_child
return root.get_data()
2019-10-05 05:14:13 +00:00
def get_left_most(root: MyNode) -> Any:
while True:
left_child = root.get_left()
if left_child is None:
break
root = left_child
return root.get_data()
2019-10-05 05:14:13 +00:00
def del_node(root: MyNode, data: Any) -> MyNode | 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)
root.set_data(temp_data)
root.set_right(del_node(right_child, temp_data))
elif left_child is not None:
root = left_child
elif right_child is not None:
root = right_child
else:
return None
elif root.get_data() > data:
if left_child is None:
print("No such data")
return root
else:
root.set_left(del_node(left_child, data))
else: # root.get_data() < data
if right_child is None:
return root
else:
root.set_right(del_node(right_child, data))
if get_height(right_child) - get_height(left_child) == 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:
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
2019-10-05 05:14:13 +00:00
2018-11-19 19:34:44 +00:00
class AVLtree:
"""
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
2019-10-05 05:14:13 +00:00
def get_height(self) -> int:
return get_height(self.root)
2019-10-05 05:14:13 +00:00
def insert(self, data: Any) -> None:
2019-10-05 05:14:13 +00:00
print("insert:" + str(data))
self.root = insert_node(self.root, data)
def del_node(self, data: Any) -> None:
2019-10-05 05:14:13 +00:00
print("delete:" + str(data))
if self.root is None:
print("Tree is empty!")
return
2019-10-05 05:14:13 +00:00
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()
2018-11-19 19:34:44 +00:00
q.push(self.root)
layer = self.get_height()
if layer == 0:
return output
2018-11-19 19:34:44 +00:00
cnt = 0
while not q.is_empty():
2018-11-19 19:34:44 +00:00
node = q.pop()
2019-10-05 05:14:13 +00:00
space = " " * int(math.pow(2, layer - 1))
output += space
2018-11-19 19:34:44 +00:00
if node is None:
output += "*"
q.push(None)
q.push(None)
2018-11-19 19:34:44 +00:00
else:
output += str(node.get_data())
q.push(node.get_left())
q.push(node.get_right())
output += space
2018-11-19 19:34:44 +00:00
cnt = cnt + 1
for i in range(100):
2019-10-05 05:14:13 +00:00
if cnt == math.pow(2, i) - 1:
layer = layer - 1
if layer == 0:
output += "\n*************************************"
return output
output += "\n"
2018-11-19 19:34:44 +00:00
break
output += "\n*************************************"
return output
2019-10-05 05:14:13 +00:00
def _test() -> None:
import doctest
doctest.testmod()
2019-10-05 05:14:13 +00:00
2018-11-19 19:34:44 +00:00
if __name__ == "__main__":
_test()
2018-11-19 19:34:44 +00:00
t = AVLtree()
lst = list(range(10))
random.shuffle(lst)
for i in lst:
t.insert(i)
print(str(t))
random.shuffle(lst)
for i in lst:
t.del_node(i)
print(str(t))