python_avl.py

import random
def main():
    a = AVLTree()
    for i in range(0, 100):
        # print '-----------------------------------'
        a.insert(random.randrange(0, 100000))
    a.pr()
    assert len(list(a.traverse())) == 100, len(list(a.traverse()))
    for v in a.traverse():
        print v,

class AVLTree(object):
    def __init__(self):
        self.root = None
    def insert(self, value):
        if self.root is None:
            self.root = A(value)
            return
        self.root = self.root.insert(A(value))
    def traverse(self):
        for node in self.root.traverse():
            yield node.value
    def pr(self):
        self.root.pr()


class A(object):
    def __init__(self, value):
        self.value = value
        self.left = None
        self.right = None
        self.height = 1
    def count(self):
        count = 1
        if self.left:
            count += self.left.count()
        if self.right:
            count += self.right.count()
        return count

    def traverse(self):
        if self.left:
            for node in self.left.traverse():
                yield node
        yield self
        if self.right:
            for node in self.right.traverse():
                yield node
    def root(self):
        pass
    def pr(self, tab_level=0):
        print '\t' * tab_level, '-', self.value, '(', self.height, ')', \
                    '(', self.balance, ')', (self.ll , self.rl)
        if self.left:
            self.left.pr(tab_level=tab_level+1)
        else:
            print '\t' * (tab_level+1), '-', 'NONE'
        if self.right:
            self.right.pr(tab_level=tab_level+1)
        else:
            print '\t' * (tab_level+1), '-', 'NONE'

    @property
    def ll(self):
        if self.left is None:
            return 0
        return self.left.height
    @property
    def rl(self):
        if self.right is None:
            return 0
        return self.right.height
    @property
    def balance(self):
        return self.ll - self.rl

    def insert(self, a):
        # print 'INS', a.value, 'into', self.value
        # Insert
        if a.value < self.value:
            if self.left:
                self.left = self.left.insert(a)
            else:
                self.left = a
        else:
            if self.right:
                self.right = self.right.insert(a)
            else:
                self.right = a
        self.height = max(self.ll, self.rl) + 1
        return self.rebalance()
    def rebalance(self):
        # Rebalance
        diff = self.ll - self.rl
        if abs(self.balance) < 2:
            return self

        assert -2 <= self.balance <= 2
        if self.balance == -2:
            rb = 0 if not self.right else self.right.balance
            assert -1 <= rb <= 1
            if rb == 1:
                self.right = self.right.right_rotation()
            res = self.left_rotation()
            return res
        elif self.balance == 2:
            lb = 0 if not self.left else self.left.balance
            assert -1 <= lb <= 1
            if lb == -1:
                self.left = self.left.left_rotation()
            return self.right_rotation()
        assert False

    def reset_height(self):
        self.height = max(self.rl, self.ll) + 1
    def left_rotation(self):
        right = self.right
        right_left = None if not right else right.left
        if not right:
            return self
        self.right = right_left
        right.left = self
        self.reset_height()
        right.reset_height()
        return right

    def right_rotation(self):
        left = self.left
        left_right = None if not left else left.right
        if not left:
            return self
        
        self.left = left_right
        left.right = self
        self.reset_height()
        left.reset_height()
        return left


main()