Python/data_structures/AVL/AVL.py
2017-10-25 01:37:11 +05:30

181 lines
4.8 KiB
Python

'''
A AVL tree
'''
class Node:
def __init__(self, label):
self.label = label
self._parent = None
self._left = None
self._right = None
self.height = 0
@property
def right(self):
return self._right
@right.setter
def right(self, node):
if node is not None:
node._parent = self
self._right = node
@property
def left(self):
return self._left
@left.setter
def left(self, node):
if node is not None:
node._parent = self
self._left = node
@property
def parent(self):
return self._parent
@parent.setter
def parent(self, node):
if node is not None:
self._parent = node
self.height = self.parent.height + 1
else:
self.height = 0
class AVL:
def __init__(self):
self.root = None
self.size = 0
def insert(self, value):
node = Node(value)
if self.root is None:
self.root = node
self.root.height = 0
self.size = 1
else:
# Same as Binary Tree
dad_node = None
curr_node = self.root
while True:
if curr_node is not None:
dad_node = curr_node
if node.label < curr_node.label:
curr_node = curr_node.left
else:
curr_node = curr_node.right
else:
node.height = dad_node.height
dad_node.height += 1
if node.label < dad_node.label:
dad_node.left = node
else:
dad_node.right = node
self.rebalance(node)
self.size += 1
break
def rebalance(self, node):
n = node
while n is not None:
height_right = n.height
height_left = n.height
if n.right is not None:
height_right = n.right.height
if n.left is not None:
height_left = n.left.height
if abs(height_left - height_right) > 1:
if height_left > height_right:
left_child = n.left
if left_child is not None:
h_right = (right_child.right.height
if (right_child.right is not None) else 0)
h_left = (right_child.left.height
if (right_child.left is not None) else 0)
if (h_left > h_right):
self.rotate_left(n)
break
else:
self.double_rotate_right(n)
break
else:
right_child = n.right
if right_child is not None:
h_right = (right_child.right.height
if (right_child.right is not None) else 0)
h_left = (right_child.left.height
if (right_child.left is not None) else 0)
if (h_left > h_right):
self.double_rotate_left(n)
break
else:
self.rotate_right(n)
break
n = n.parent
def rotate_left(self, node):
aux = node.parent.label
node.parent.label = node.label
node.parent.right = Node(aux)
node.parent.right.height = node.parent.height + 1
node.parent.left = node.right
def rotate_right(self, node):
aux = node.parent.label
node.parent.label = node.label
node.parent.left = Node(aux)
node.parent.left.height = node.parent.height + 1
node.parent.right = node.right
def double_rotate_left(self, node):
self.rotate_right(node.getRight().getRight())
self.rotate_left(node)
def double_rotate_right(self, node):
self.rotate_left(node.getLeft().getLeft())
self.rotate_right(node)
def empty(self):
if self.root is None:
return True
return False
def preShow(self, curr_node):
if curr_node is not None:
self.preShow(curr_node.left)
print(curr_node.label, end=" ")
self.preShow(curr_node.right)
def preorder(self, curr_node):
if curr_node is not None:
self.preShow(curr_node.left)
self.preShow(curr_node.right)
print(curr_node.label, end=" ")
def getRoot(self):
return self.root
t = AVL()
t.insert(1)
t.insert(2)
t.insert(3)
# t.preShow(t.root)
# print("\n")
# t.insert(4)
# t.insert(5)
# t.preShow(t.root)
# t.preorden(t.root)