Recitation04 Notes

From 6.006 Wiki

Revision as of 20:46, 17 September 2008

BST / AVL implementation

• Binary Search Tree

# 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):
    """ 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)