Binary Search Tree in Python

binary_search_tree.py

#!/usr/bin/env python

"""
A binary search tree implementation, from:

"Python Algorithms: Mastering Basic Algorithms in the Python Language"

by  	Magnus Lie Hetland
ISBN: 	9781484200551 

"""

class Node:
    lft = None
    rgt = None
    def __init__(self, key, val):
        self.key = key
        self.val = val

def insert(node, key, val):
    if node is None:
        return Node(key, val) # empty leaf: add node here
    if node.key == key:
        node.val = val # found the key: replace val
    elif key < node.key:
        node.lft = insert(node.lft, key, val) # go left
    else:
        node.rgt = insert(node.rgt, key, val) # go right
    return node

def search(node, key):
    if node is None:
        raise KeyError # empty leaf: it's not here
    if node.key == key:
        return node.val # found key: return val
    elif key < node.key:
        return search(node.lft, key) # go left
    else:
        return search(node.rgt, key) # go right

def traverse_recursively(node):
    """Perform an in-order traversal recursively"""
    
    nodes = []
    
    if node is not None:
        nodes.extend(traverse_recursively(node.lft))
        nodes.append(node)
        nodes.extend(traverse_recursively(node.rgt))
        
    return nodes

def traverse_stack(node):
    """Perform an in-order traversal iteratively using a stack"""
    
    nodes = []
    
    stack = []
    while True:
        while node is not None:
            stack.append(node)
            node = node.lft

        if len(stack) == 0:
            break

        node = stack.pop()
        nodes.append(node)
        node = node.rgt
        
    return nodes

class Tree:
    # a simple wrapper
    root = None
    def __setitem__(self, key, val):
        self.root = insert(self.root, key, val)
    def __getitem__(self, key):
        return search(self.root, key)
    def __contains__(self, key):
        try:
            search(self.root, key)
        except KeyError:
            return False
        return True
    def rtraverse(self):
        return traverse_recursively(self.root)
    def straverse(self):
        return traverse_stack(self.root)

Usage:

tree = Tree()
tree[0] = "a"
tree[1] = "b"

0 in tree
True

1 in tree
True

2 in tree
False

Usage with new in-order traversal functions:

tree = Tree()
tree[14] = 14 # key drives the sequence ordering, val can be any arbitrary data
tree[35] = 35
tree[10] = 10
tree[19] = 19
tree[31] = 31
tree[42] = 42

# traverse in key order recursively
map(lambda x : x.val, tree.rtraverse())
[10, 14, 19, 31, 35, 42]

# traverse in key order with a stack
map(lambda x : x.val, tree.straverse())
[10, 14, 19, 31, 35, 42]