A binary search tree in Python

As part of my series on concepts you should never, ever need to reimplement in a real world scenario, here is a binary search tree in Python. This particular example contains a bit of duplicate code, but I doubt the benefits of changing this would outweigh the loss in readability, especially for a learning example.

For the sake of doing things correctly, we extend collections.Container, a Python abstract class shared by most other collections.

from collections import Container
class Node(object):
    left = None
    right = None
    value = None

class BinarySearchTree(Container):

    root = Node()

    # Allows "if item in tree:"
    def __contains__(self, item):
        return self._contains(item, self.root)

    def _contains(self, item, node):
        if item > node.value:
            if node.right:
                return self._contains(item, node.right)
            else:
                return False

        elif item < node.value:
            if node.left:
                return self._contains(item, node.left)
            else:
                return False

        else:
            return True

    def add(self, item):
        self._add(item, self.root)

    def _add(self, item, node):
        if node.value:
            if item > node.value:
                if node.right:
                    self._add(item, node.right)
                else:
                    node.right = Node()
                    node.right.value = item
            elif item < node.value:
                if node.left:
                    self._add(item, node.left)
                else:
                    node.left = Node()
                    node.left.value = item
            else:  # Node does not exist or has same value as item
                node = Node()
                node.value = item
        else:
            node.value = item

    # Useful for debugging. Prints the whole tree
    def __repr__(self):
        return self._repr(self.root)

    def _repr(self, node, indent=0):
        string = ('\t' * indent) + str(node.value) + '\n'

        if node.left:
            string += ('\t' * (indent+1)) + 'Left:\n'
            string += self._repr(node.left, indent+1)
        if node.right:
            string += ('\t' * (indent+1)) + 'Right:\n'
            string += self._repr(node.right, indent+1)

        return string

And now we can test our tree:

from random import randint

# Add random values
values = []
tree = BinarySearchTree()
for i in range(40):
    val = randint(0,100)
    tree.add(val)
    values.append(val)

print(tree)  # uses __repr__
print(values)
print(25 in tree)  # uses __contains__
print(25 in values)

Leave a Reply

Your email address will not be published. Required fields are marked *

To create code blocks or other preformatted text, indent by four spaces:

    This will be displayed in a monospaced font. The first four 
    spaces will be stripped off, but all other whitespace
    will be preserved.
    
    Markdown is turned off in code blocks:
     [This is not a link](http://example.com)

To create not a block, but an inline code span, use backticks:

Here is some inline `code`.

For more help see http://daringfireball.net/projects/markdown/syntax