Python/data_structures/binary_tree/avl_tree.py

289 lines
7.5 KiB
Python
Raw Normal View History

2019-10-05 05:14:13 +00:00
"""
An auto-balanced binary tree!
2019-10-05 05:14:13 +00:00
"""
2018-11-19 19:34:44 +00:00
import math
import random
2019-10-05 05:14:13 +00:00
2018-11-19 19:34:44 +00:00
class my_queue:
def __init__(self):
self.data = []
self.head = 0
self.tail = 0
2019-10-05 05:14:13 +00:00
2018-11-19 19:34:44 +00:00
def isEmpty(self):
return self.head == self.tail
2019-10-05 05:14:13 +00:00
def push(self, data):
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
2018-11-19 19:34:44 +00:00
def pop(self):
ret = self.data[self.head]
self.head = self.head + 1
return ret
2019-10-05 05:14:13 +00:00
2018-11-19 19:34:44 +00:00
def count(self):
return self.tail - self.head
2019-10-05 05:14:13 +00:00
2018-11-19 19:34:44 +00:00
def print(self):
print(self.data)
print("**************")
2019-10-05 05:14:13 +00:00
print(self.data[self.head : self.tail])
2018-11-19 19:34:44 +00:00
class my_node:
2019-10-05 05:14:13 +00:00
def __init__(self, data):
2018-11-19 19:34:44 +00:00
self.data = data
self.left = None
self.right = None
self.height = 1
2019-10-05 05:14:13 +00:00
2018-11-19 19:34:44 +00:00
def getdata(self):
return self.data
2019-10-05 05:14:13 +00:00
2018-11-19 19:34:44 +00:00
def getleft(self):
return self.left
2019-10-05 05:14:13 +00:00
2018-11-19 19:34:44 +00:00
def getright(self):
return self.right
2019-10-05 05:14:13 +00:00
2018-11-19 19:34:44 +00:00
def getheight(self):
return self.height
2019-10-05 05:14:13 +00:00
def setdata(self, data):
2018-11-19 19:34:44 +00:00
self.data = data
return
2019-10-05 05:14:13 +00:00
def setleft(self, node):
2018-11-19 19:34:44 +00:00
self.left = node
return
2019-10-05 05:14:13 +00:00
def setright(self, node):
2018-11-19 19:34:44 +00:00
self.right = node
return
2019-10-05 05:14:13 +00:00
def setheight(self, height):
2018-11-19 19:34:44 +00:00
self.height = height
return
2019-10-05 05:14:13 +00:00
2018-11-19 19:34:44 +00:00
def getheight(node):
if node is None:
return 0
return node.getheight()
2019-10-05 05:14:13 +00:00
def my_max(a, b):
2018-11-19 19:34:44 +00:00
if a > b:
return a
return b
def leftrotation(node):
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.getdata())
2018-11-19 19:34:44 +00:00
ret = node.getleft()
node.setleft(ret.getright())
ret.setright(node)
2019-10-05 05:14:13 +00:00
h1 = my_max(getheight(node.getright()), getheight(node.getleft())) + 1
2018-11-19 19:34:44 +00:00
node.setheight(h1)
2019-10-05 05:14:13 +00:00
h2 = my_max(getheight(ret.getright()), getheight(ret.getleft())) + 1
2018-11-19 19:34:44 +00:00
ret.setheight(h2)
return ret
2019-10-05 05:14:13 +00:00
2018-11-19 19:34:44 +00:00
def rightrotation(node):
2019-10-05 05:14:13 +00:00
"""
2018-11-22 06:33:50 +00:00
a mirror symmetry rotation of the leftrotation
2019-10-05 05:14:13 +00:00
"""
print("right rotation node:", node.getdata())
2018-11-19 19:34:44 +00:00
ret = node.getright()
node.setright(ret.getleft())
ret.setleft(node)
2019-10-05 05:14:13 +00:00
h1 = my_max(getheight(node.getright()), getheight(node.getleft())) + 1
2018-11-19 19:34:44 +00:00
node.setheight(h1)
2019-10-05 05:14:13 +00:00
h2 = my_max(getheight(ret.getright()), getheight(ret.getleft())) + 1
2018-11-19 19:34:44 +00:00
ret.setheight(h2)
return ret
2019-10-05 05:14:13 +00:00
2018-11-22 06:33:50 +00:00
def rlrotation(node):
2019-10-05 05:14:13 +00:00
r"""
A A Br
2018-11-22 06:33:50 +00:00
/ \ / \ / \
B C RR Br C LR B A
/ \ --> / \ --> / / \
Bl Br B UB Bl UB C
2018-11-22 06:33:50 +00:00
\ /
UB Bl
RR = rightrotation LR = leftrotation
2019-10-05 05:14:13 +00:00
"""
2018-11-22 06:33:50 +00:00
node.setleft(rightrotation(node.getleft()))
return leftrotation(node)
2019-10-05 05:14:13 +00:00
2018-11-19 19:34:44 +00:00
def lrrotation(node):
node.setright(leftrotation(node.getright()))
return rightrotation(node)
2019-10-05 05:14:13 +00:00
def insert_node(node, data):
2018-11-19 19:34:44 +00:00
if node is None:
return my_node(data)
if data < node.getdata():
2019-10-05 05:14:13 +00:00
node.setleft(insert_node(node.getleft(), data))
if (
getheight(node.getleft()) - getheight(node.getright()) == 2
): # an unbalance detected
if (
data < node.getleft().getdata()
): # new node is the left child of the left child
2018-11-19 19:34:44 +00:00
node = leftrotation(node)
else:
2019-10-05 05:14:13 +00:00
node = rlrotation(node) # new node is the right child of the left child
2018-11-19 19:34:44 +00:00
else:
2019-10-05 05:14:13 +00:00
node.setright(insert_node(node.getright(), data))
2018-11-19 19:34:44 +00:00
if getheight(node.getright()) - getheight(node.getleft()) == 2:
if data < node.getright().getdata():
node = lrrotation(node)
else:
node = rightrotation(node)
2019-10-05 05:14:13 +00:00
h1 = my_max(getheight(node.getright()), getheight(node.getleft())) + 1
2018-11-19 19:34:44 +00:00
node.setheight(h1)
return node
2019-10-05 05:14:13 +00:00
def getRightMost(root):
while root.getright() is not None:
root = root.getright()
return root.getdata()
2019-10-05 05:14:13 +00:00
def getLeftMost(root):
while root.getleft() is not None:
root = root.getleft()
return root.getdata()
2019-10-05 05:14:13 +00:00
def del_node(root, data):
if root.getdata() == data:
if root.getleft() is not None and root.getright() is not None:
temp_data = getLeftMost(root.getright())
root.setdata(temp_data)
2019-10-05 05:14:13 +00:00
root.setright(del_node(root.getright(), temp_data))
elif root.getleft() is not None:
root = root.getleft()
else:
root = root.getright()
elif root.getdata() > data:
if root.getleft() is None:
print("No such data")
return root
else:
2019-10-05 05:14:13 +00:00
root.setleft(del_node(root.getleft(), data))
elif root.getdata() < data:
if root.getright() is None:
return root
else:
2019-10-05 05:14:13 +00:00
root.setright(del_node(root.getright(), data))
if root is None:
return root
if getheight(root.getright()) - getheight(root.getleft()) == 2:
if getheight(root.getright().getright()) > getheight(root.getright().getleft()):
root = rightrotation(root)
else:
root = lrrotation(root)
elif getheight(root.getright()) - getheight(root.getleft()) == -2:
if getheight(root.getleft().getleft()) > getheight(root.getleft().getright()):
root = leftrotation(root)
else:
root = rlrotation(root)
2019-10-05 05:14:13 +00:00
height = my_max(getheight(root.getright()), getheight(root.getleft())) + 1
root.setheight(height)
return root
2019-10-05 05:14:13 +00:00
2018-11-19 19:34:44 +00:00
class AVLtree:
def __init__(self):
self.root = None
2019-10-05 05:14:13 +00:00
2018-11-19 19:34:44 +00:00
def getheight(self):
2019-10-05 05:14:13 +00:00
# print("yyy")
2018-11-19 19:34:44 +00:00
return getheight(self.root)
2019-10-05 05:14:13 +00:00
def insert(self, data):
print("insert:" + str(data))
self.root = insert_node(self.root, data)
def del_node(self, data):
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 traversale(self): # a level traversale, gives a more intuitive look on the tree
2018-11-19 19:34:44 +00:00
q = my_queue()
q.push(self.root)
layer = self.getheight()
if layer == 0:
return
2018-11-19 19:34:44 +00:00
cnt = 0
while not q.isEmpty():
node = q.pop()
2019-10-05 05:14:13 +00:00
space = " " * int(math.pow(2, layer - 1))
print(space, end="")
2018-11-19 19:34:44 +00:00
if node is None:
2019-10-05 05:14:13 +00:00
print("*", end="")
q.push(None)
q.push(None)
2018-11-19 19:34:44 +00:00
else:
2019-10-05 05:14:13 +00:00
print(node.getdata(), end="")
2018-11-19 19:34:44 +00:00
q.push(node.getleft())
q.push(node.getright())
2019-10-05 05:14:13 +00:00
print(space, end="")
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:
print()
print("*************************************")
return
2018-11-19 19:34:44 +00:00
print()
break
print()
print("*************************************")
return
2019-10-05 05:14:13 +00:00
2018-11-19 19:34:44 +00:00
def test(self):
getheight(None)
print("****")
self.getheight()
2019-10-05 05:14:13 +00:00
2018-11-19 19:34:44 +00:00
if __name__ == "__main__":
t = AVLtree()
t.traversale()
lst = list(range(10))
random.shuffle(lst)
for i in lst:
t.insert(i)
t.traversale()
2019-10-05 05:14:13 +00:00
random.shuffle(lst)
for i in lst:
t.del_node(i)
t.traversale()