import random

class TreeNode(object):

    '''A single node in our tree.'''
    
    def __init__(self, value):
        '''Create a new node containing value.'''
        self.value = value
        self.left = None
        self.right = None
        
    def in_order(self):
        '''An inorder traversal of the subtree rooted at self.'''
        # Permuting these lines can turn this into pre/post order.
        if self.left is not None:
            self.left.in_order()
        print self.value,
        if self.right is not None:
            self.right.in_order()
            
    def height(self):
        '''Perform a post-order traversal to find the height of the subtree
        rooted at self.'''
        left_height = 0
        if self.left is not None:
            left_height = 1 + self.left.height()
        right_height = 0
        if self.right is not None:
            right_height = 1 + self.right.height()
        return max(left_height, right_height)
        
class BSTree(object):

    '''A Binary Search Tree which supports searching, insert, and removal.'''
    
    def __init__(self):
        '''Create an empty tree.'''
        self.root = None
        
    def _search(self, node, key):
        '''Recursive helper method for search.'''
        # Try and implement this iteratively.
        if node is None:
            return False
        if node.value == key:
            return True
        if node.value < key:
            return self._search(node.right, key)
        else:
            return self._search(node.left, key)
            
    def search(self, key):
        '''Return True iff key is in the tree.  Uses the recursive _search
        helper method.'''
        return self._search(self.root, key)

    def _find_and_remove_successor(self, node):
        '''Given a node with two children, find its successor, remove it from
        the tree and remove it from the tree and return it.'''
        succ_node = node.right
        succ_parent = node
        while succ_node.left is not None:
            succ_parent = succ_node
            succ_node = succ_node.left
        # succ_node is the successor of node.  Provably, this node has at most
        # one child (a right one).  To remove it, we set its parents pointer to
        # it to be its right child, and then return it.
        if succ_parent.right is succ_node:
            succ_parent.right = succ_node.right
        else:
            succ_parent.left = succ_node.right
        return succ_node
    
    def remove(self, key):
        '''If some node in the tree has value == key, remove it.'''
        if self.root is None:
            # key is not in the tree.
            return
        # If key is in the tree, find the node which contains it, and the
        # parent of that node.
        key_parent = None
        key_node = self.root
        while key_node.value != key:
            if key_node.value < key:
                key_parent = key_node
                key_node = key_node.right
            else:
                key_parent = key_node
                key_node = key_node.left
            if key_node is None:
                # key is not in the tree.
                return
        # key_node is now the node which contains key, and key_parent is its
        # parent (or None if it is the root).
        if key_node.left is not None and key_node.right is not None:
            # The node has two children (hard case).
            succ_node = self._find_and_remove_successor(key_node)
            # succ_node adopts the children of key_node.
            succ_node.left = key_node.left
            succ_node.right = key_node.right
            # Now just need to update the children of key_parent.
            if key_parent is None:
                self.root = succ_node
            else:
                if key_node is key_parent.left:
                    key_parent.left = succ_node
                else:
                    key_parent.right = succ_node
        else:
            # The node has at most one child, so all we need to do is update
            # the parent.
            child = key_node.right
            if key_node.left is not None:
                child = key_node.left
            if key_parent is None:
                # The new root is going to be child
                self.root = child
            else:
                if key_node is key_parent.left:
                    key_parent.left = child
                else:
                    key_parent.right = child
    
    def insert(self, key):
        '''Insert key into the tree at a leaf.'''
        if self.root is None:
            self.root = TreeNode(key)
        else:
            done = False
            current = self.root
            while not done:
                if current.value == key:
                    # key is already in the tree.
                    done = True
                elif current.value < key:
                    if current.right is None:
                        # This is where we want to insert key.
                        current.right = TreeNode(key)
                        Done = True
                    else:
                        # Insert into the right subtree.
                        current = current.right
                else:
                    if current.left is None:
                        # This is where we want to insert key.
                        current.left = TreeNode(key)
                        Done = True
                    else:
                        # Insert into the left subtree.
                        current = current.left
        
        
if __name__ == '__main__':
    items = range(3)
    random.shuffle(items)
    tree = BSTree()
    for x in items:
        tree.insert(x)
    print 'The height of the randomly built tree with 10 nodes is',
    print tree.root.height()
    tree.root.in_order()
    print
    random.shuffle(items)
    print items
    for i in items:
        tree.remove(i)
        if tree.root:
            tree.root.in_order()
            print


