From 6.006 Wiki
BST / AVL implementation
# Recitation04- WF3
#
# An implementation of an AVL tree using Python objects
#
import r03_BSTobjects2 as bst
def height(node):
""" Returns the height of a node """
if node is None:
return -1
else:
return node.height
def update_height(node):
""" Determines and sets the height of the node """
node.height = max(height(node.left), height(node.right)) + 1
class AVL(rec03.BST):
"""
AVL binary search tree implementation.
"""
def left_rotate(self, x):
y = x.right #Step 1
y.parent = x.parent
if y.parent is None:
self.root = y
else:
if y.parent.left is x:
y.parent.left = y
elif y.parent.right is x:
y.parent.right = y
x.right = y.left #Step 2
if x.right is not None:
x.right.parent = x
y.left = x #Step 3
x.parent = y
update_height(x)
update_height(y)
def right_rotate(self, x):
y = x.left #Step 1
y.parent = x.parent
if y.parent is None:
self.root = y
else:
if y.parent.left is x:
y.parent.left = y
elif y.parent.right is x:
y.parent.right = y
x.left = y.right #Step 2
if x.left is not None:
x.left.parent = x
y.right = x #Step 3
x.parent = y
update_height(x)
update_height(y)
def insert(self, t):
"""Insert key t into this tree"""
node = bst.BST.insert(self, t)
self.rebalance(node)
def delete(self, t):
node = bst.BST.delete(self,t)
self.rebalance(node.parent)
def rebalance(self, node):
while node is not None:
update_height(node)
if height(node.left) >= 2 + height(node.right):
if height(node.left.left) >= height(node.left.right):
self.right_rotate(node)
else:
self.left_rotate(node.left)
self.right_rotate(node)
elif height(node.right) >= 2 + height(node.left):
if height(node.right.right) >= height(node.right.left):
self.left_rotate(node)
else:
self.right_rotate(node.right)
self.left_rotate(node)
node = node.parent
###############
if __name__ == '__main__':
L = [10, 5, 8, 7, 4, 11, 1, 3, 2, 9, 6, 12]
A = AVL()
B = bst.BST()
for item in L:
A.insert(item)
B.insert(item)