avlmap.py

from avltree import AVLTree

class AVLMap(object):
    def __init__(self):
        self._store = AVLTree()

    def __setitem__(self, key, value):
        self._store.insert((key, value))

    def __getitem__(self, key):
        node = self._store.search(key, lambda i: i[0])
        if node.value[0] != key:
            raise KeyError
        return node.value[1]

    def __delitem__(self, key):
        node = self._store.search(key, lambda i: i[0])
        if node.value[0] != key:
            raise KeyError
        node.delete(node.value)

    def __repr__(self):
        items = []
        for k, v in self._store:
            items.append(repr(k) + ': ' + repr(v))
        return '{' + ', '.join(items) + '}'

if __name__ == '__main__':
    d = AVLMap()
    print d
    d['foo'] = 'bar'
    print d
    print d['foo']
    d['baz'] = 'quux'
    print d
    print d['foo']
    del d['foo']
    print d

avltree.py

class Empty:
    pass

class AVLTree(object):
    value = None
    left = None
    right = None

    def __init__(self, value=Empty, left=None, right=None, parent=None):
        if value is Empty and parent:
            raise ValueError("Cannot use special value Empty")
        self.value = value
        self.left = left
        self.right = right
        self.parent = parent

    @property
    def height(self):
        if self.left:
            lh = self.left.height
        else:
            lh = 0
        if self.right:
            rh = self.right.height
        else:
            rh = 0
        return max(lh, rh) + 1

    def search(self, value, mutator=None):
        """Return the sub-tree which contains value value"""
        if mutator is None:
            mutator = lambda x: x
        if value == mutator(self.value):
            return self
        if value < self.value:
            if self.left:
                return self.left.search(value)
            return self  # self contains -infinity to right, ie value
        if value > self.value:
            if self.right:
                return self.right.search(value)
            return self

    def insert(self, value):
        if value is Empty:
            raise ValueError("Cannot insert special value Empty")
        if self.value is Empty:
            self.value = value
            return
        if value <= self.value:
            if self.left:
                return self.left.insert(value)
            else:
                self.left = AVLTree(value, parent=self)
            self.rebalance()
            return
        if self.right:
            self.right.insert(value)
        else:
            self.right = AVLTree(value, parent=self)
        self.rebalance()

    def insert_tree(self, other):
        for value in other:
            self.insert(value)

    def __iter__(self):
        if self.left:
            for value in self.left:
                yield value
        if self.value is not Empty:
            yield self.value
        if self.right:
            for value in self.right:
                yield value

    def _replace_with(self, other):
        self.left = other.left
        self.value = other.value
        self.right = other.right

    def remove_child(self, other):
        if other is self.left:
            self.left = None
        elif other is self.right:
            self.right = None
        else:
            raise ValueError('whut')

    def delete(self, value):
        if value != self.value:
            tree = self.search(value)
            if tree.value != value:
                raise KeyError("Value not in tree")
            return tree.delete(value)

        if not self.left and not self.right:
            if not self.parent:
                self.value = Empty
                return
            self.parent.remove_child(self)
            return

        if not self.left:
            return self._replace_with(self.right)
        if not self.right:
            return self._replace_with(self.left)

        right_min = self.right.min_node()
        self.value = right_min.value
        # The following is basically right_min.delete(value)
        if right_min.right:
            right_min.value = right_min.right.value
            right_min.right = None
        elif right_min.parent is not self:
            right_min.parent.left = None
        else:  # We've deleted the last remaining right node
            self.right = None
        self.rebalance()

    def rebalance(self):
        bf = self.balance_factor()
        dh = abs(bf)
        if dh <= 1:
            if self.parent:
                return self.parent.rebalance()
            return

        if bf > 0:  # Left heavy
            left_bf = self.left.balance_factor()
            if left_bf >= 0:
                self.rotate_right()
                return
            else:
                self.left.rotate_left()
                self.rotate_right()
        else:  # right heavy
            right_bf = self.right.balance_factor()
            if right_bf < 0:
                self.rotate_left()
                return
            else:
                self.right.rotate_right()
                self.rotate_left()

        if self.parent:
            return self.parent.rebalance()
        return

    def rotate_right(self):
        if not self.left:
            raise ValueError("Cannot rotate right")
        a, b, g = self.left.left, self.left.right, self.right
        self.right = AVLTree(self.value, left=b, right=g, parent=self)
        if b:
            b.parent = self.right
        if g:
            g.parent = self.right
        self.value = self.left.value
        self.left = a
        if self.left:
            self.left.parent = self

    def rotate_left(self):
        if not self.right:
            raise ValueError("Cannot rotate left")
        a, b, g = self.left, self.right.left, self.right.right
        self.left = AVLTree(self.value, left=a, right=b, parent=self)
        if a:
            a.parent = self.left
        if b:
            b.parent = self.left

        self.value = self.right.value
        self.right = g
        if self.right:
            self.right.parent = self


    def printtree(self, level=0):
        if self.right:
            self.right.printtree(level + 1)
        print((" " * level) + str(self.value))
        if self.left:
            self.left.printtree(level + 1)

    def min_node(self):
        if not self.left:
            return self
        return self.left.min_node()

    def min(self):
        return self.min_node().value

    def max_node(self):
        if not self.right:
            return self
        return self.right.max_node()

    def max(self):
        return self.max_node().value

    def balance_factor(self):
        if not self.left:
            lh = 0
        else:
            lh = self.left.height

        if not self.right:
            rh = 0
        else:
            rh = self.right.height

        return lh - rh


if __name__ == '__main__':
    t = AVLTree(5)
    t.printtree()
    print
    t.insert(3)
    t.insert(6)
    t.printtree()
    print

    t.insert(4)
    t.insert(2)
    t.insert(1)
    t.insert(0)
    t.insert(-1)
    t.printtree()
    print

    t.delete(-1)
    t.delete(1)
    t.printtree()
    print
    t.delete(5)
    t.printtree()
    print
    t.delete(4)
    t.printtree()
    print

    t.delete(0)
    t.delete(2)
    t.delete(3)
    t.printtree()
    t.delete(6)
    t.printtree()