mirror of
https://github.com/TheAlgorithms/Python.git
synced 2024-11-24 13:31:07 +00:00
181 lines
4.8 KiB
Python
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)
|