binary_tree.py

#!/usr/bin/env python
# encoding: utf-8


class Node(object):

    def __init__(self, value, left=None, right=None):

        self.value = value
        self.left = left
        self.right = right

    def __str__(self):
        return str(self.value)


class BinarySearchTree(object):

    def __init__(self):

        self.root = None

    def insert(self, value):

        if not self.root:
            self.root = Node(value)

        node = self.root

        while node:

            if node.value == value:
                return

            elif node.value < value:

                if not node.right:
                    node.right = Node(value)

                node = node.right

            else:

                if not node.left:
                    node.left = Node(value)

                node = node.left

    def contains(self, value):

        node = self.root

        while node:

            if node.value == value:
                return True

            if node.value > value:
                node = node.left
            else:
                node = node.right

        return False

    def get_min(self):

        if not self.root:
            return None

        node = self.root

        while node.left:
            node = node.left

        return node.value

    def get_max(self):

        if not self.root:
            return None

        node = self.root

        while node.right:

            node = node.right

        return node.value

    def remove(self, value):
        self.root = self.__remove(value, self.root)

    def pop_min(self):

        if not self.root:
            return None

        if self.root.left is None:
            value = self.root.value
            self.root = self.right
            return value

        return self.__remove_min(self.root)

    def __remove(self, value, node):

        if node is None:
            return None

        if node.value < value:
            node.right = self.__remove(value, node.right)
        elif node.value > value:
            node.left = self.__remove(value, node.left)
        else:
            if node.left and node.right:
                if node.right.left is None:
                    node.value = node.right.value
                    node.right = node.right.right
                else:
                    node.value = self.__remove_min(node.right)
            else:
                node = node.left if node.left else node.right

        return node

    def __remove_min(self, node):

        if node is None:
            return None

        if node.left is None:
            return node.value

        while node and node.left and node.left.left:
            node = node.left

        result = node.left.value

        node.left = node.left.right

        return result

    def level_print(self, ident='', node=None, last=True):

        node = node if node else self.root

        if not node:
            return

        if last:
            print ident, '\-' + str(node)
            ident += '  '
        else:
            print ident, '|-' + str(node)
            ident += ' |'

        if node.left:
            self.level_print(ident, node.left, last=False)

        if node.right:
            self.level_print(ident, node.right)

tree = BinarySearchTree()

tree.insert(20)
tree.insert(8)
tree.insert(10)
tree.insert(333)
tree.insert(5)
tree.insert(7)
tree.insert(9)
tree.insert(200)
tree.insert(100)
tree.insert(400)
tree.insert(600)
tree.insert(500)
tree.insert(1)
tree.insert(2)

tree.level_print()

print tree.pop_min()

tree.level_print()

print tree.pop_min()