trees.py

import random

class TreeNode:
    def __init__(self, key, value, left=None, right=None):
        self.key = key
        self.value = value
        self.left = left
        self.right = right

    def __repr__(self):
        def __str_none(x):
            if x == None:
                return ""
            return str(x)

        return "[%s %s(%s) %s]" % \
               (__str_none(self.left),
                self.key,
                self.value,
                __str_none(self.right))

class BinarySearchTree:
    def __init__(self):
        self.root = None

    def __repr__(self):
        if self.root == None:
            return "BinarySearchTree []"
        return "BinarySearchTree %s" % self.root

    def lookup(self, key):
        def __lookup(node, key):
            if node == None:
                return None

            if key == node.key:
                return node.value
            if key < node.key:
                return __lookup(node.left, key)
            return __lookup(node.right, key)
        
        return __lookup(self.root, key)

    def insert(self, key, value=None):
        def __insert(node, key, value):
            if node == None:
                return TreeNode(key, value)
            
            if key == node.key:
                node.value = value
            elif key < node.key:
                node.left = __insert(node.left, key, value)
            else:
                node.right = __insert(node.right, key, value)
            return node
        
        self.root = __insert(self.root, key, value)

    def size(self):
        def __size(node):
            if node == None:
                return 0
            return 1 + __size(node.left) + __size(node.right)
        
        return __size(self.root)

    def depth(self):
        def __depth(node):
            if node == None:
                return -1
            return 1 + max(__depth(node.left), __depth(node.right))
        
        return __depth(self.root)

    def min(self):
        if self.root == None:
            return None
        node = self.root
        while node.left != None:
            node = node.left
        return node.key

    def pretty_print(self, format_string=None):
        if format_string == None:
            return pretty_print(self.root)
        return pretty_print(self.root, format_string)

    def __iter__(self):
        return self.traverse()

    def traverse(self):
        def __traverse(node):
            if node.left != None:
                for elem in __traverse(node.left):
                    yield elem
            yield node.key, node.value
            if node.right != None:
                for elem in __traverse(node.right):
                    yield elem

        return __traverse(self.root)

def pretty_print(node, format_string="{key}({value})"):
    def __get(lst, i):
        if i < len(lst):
            return lst[i]
        return ""

    def __pretty_print(node):
        if node == None:
            return []
        current = format_string.format(**node.__dict__)
        left = __pretty_print(node.left)
        right = __pretty_print(node.right)
        left_width = max([len(x) for x in left] + [0])
        right_width = max([len(x) for x in right] + [0])
        current_width = 2 * (len(current) // 2) + 1

        out = [" " * left_width + current + " " * right_width]
        for i in range(max(len(left), len(right))):
            line = "{0:^{1}}".format(__get(left, i), left_width) + \
                   "{0:^{1}}".format("|", current_width) + \
                   "{0:^{1}}".format(__get(right, i), right_width)
            out += [line]
        return out

    return "\n".join(__pretty_print(node))

def pretty_print_keys(t):
    print(t.pretty_print("{key:^3}"))

def build_tree(lst):
    t = BinarySearchTree()
    for elem in lst:
        if isinstance(elem, tuple):
            t.insert(*elem)
        else:
            t.insert(elem, elem)
    return t

def check_depth(n):
    tree = BinarySearchTree()
    for i in range(n):
        tree.insert(random.randrange(1, n**2))
    return tree.depth()

def check_many_dpeths(depths, attempts=10):
    for d in depths:
        tests = [check_depth(d) for i in range(attempts)]
        print((d, sum(tests) / attempts, tests))