Binary Search Tree

binary_search_tree.py

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

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

    def __len__(self):
        return self.size

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

    def find(self, key):
        if self.root:
            return self._find(key, self.root)
        return None

    def _find(self, key, current):
        if not current:
            return None
        elif current.key == key:
            return current
        elif key < current.key:
            if current.left:
                return self._find(key, current.left)
            return current
        else:
            if current.right:
                return self._find(key, current.right)
            return current 

    def insert(self, key):
        if self.root:
            return self._insert(key, self.root)
        else:
            self.root = TreeNode(key)
        self.size += 1

    def _insert(self, key, current):