Python Binary Search Tree

As an exercise to get better in Python, I decided to implement a Binary Search Tree. This implementation uses two classes: BSTNode and BSTree.

Because most methods were implemented using recursion, I created a public method, for most of them, that wraps a private one, which is in charge of making the real work. I did that, because I didn't wanted to give the user the responsibility of passing the root of the tree to every method call. With that, the user can write my_tree.insert(10) instead of my_tree.insert(my_tree.root, 10).

Is this a reasonable way of solve this problem? Is there a better one? Are private methods and attributes commonly used in Python?

The implementation:

class BSTNode(object):
""" Node class used by the Binary Search Tree. """
def __init__(self, value):
self.value = value
self.size = 1
self.left = None
self.right = None

def compute_size(self):
""" Computes the self size according to its children sizes. """
self.size = 1
if self.left:
self.size = self.size + self.left.size
if self.right:
self.size = self.size + self.right.size

class BSTree(object):
"""
Binary Search Tree implementation which doesn't allow repeated
Nodes.
"""
def __init__(self):
self.__root = None
self.__size = 0

###############
# Public API  #
###############
def pre_order_traversal(self, fn=None):
"""
Traverse tree in pre-order and apply fn to every Node value.
"""
if fn is None:
def fn(x): return print(x)
self.__pre_order_traversal(self.__root, fn)

def in_order_traversal(self, fn=None):
"""
Traverse tree in in-order and apply fn to every Node value.
"""
if fn is None:
def fn(x): return print(x)
self.__in_order_traversal(self.__root, fn)

def post_order_traversal(self, fn=None):
"""
Traverse tree in post-order and apply fn to every Node value.
"""
if fn is None:
def fn(x): return

print(x)
self.__pre_order_traversal(self.__root, fn)

def insert(self, value):
""" Inserts a Node. Doesn't insert duplicated values."""
if self.__root is None:
self.__root = BSTNode(value)
else:
self.__root = self.__insert(self.__root, value)
self.__size = self.__root.size

def remove(self, value):
"""
Removes a Node which contains the value value.
To remove a Node, three cases must be handled.
Case 1: leaf node
-> delete it
Case 2: node has one child
-> delete node and put its child in its place
Case 3: node has two children
-> delete node and put its smallest child from its right branch in its place
"""
if self.__root:
self.__root = self.__remove(self.__root, value)

def contains(self, value):
""" Returns True if value is found. """
return self.__contains(self.__root, value)

def size(self):
""" Returns the number of elements inside the BST. """
return self.__size

###############
# Private API #
###############
def __pre_order_traversal(self, node, fn):
if node is None:
return
fn(node.value)
if node.left:
self.__pre_order_traversal(node.left, fn)
if node.right:
self.__pre_order_traversal(node.right, fn)

def __in_order_traversal(self, node, fn):
if node is None:
return
if node.left:
self.__in_order_traversal(node.left, fn)
fn(node.value)
if node.right:
self.__in_order_traversal(node.right, fn)

def __post_order_traversal(self, node, fn):
if node is None:
return
if node.left:
self.__post_order_traversal(node.left, fn)
if node.right:
self.__post_order_traversal(node.right, fn)
fn(node.value)

def __insert(self, node, value):
if node is None:
return BSTNode(value)
if node.value > value:
node.left = self.__insert(node.left, value)
elif node.value < value:
node.right = self.__insert(node.right, value)
node.compute_size()
return node

def __remove(self, node, value):
if node.value == value:

# Case 1
if node.left is None and node.right is None:
return None

# Case 2
elif node.left and node.right is None:
return node.left

# Case 2
elif node.left is None and node.right:
return node.right

# Case 3
else:
parent_node = node
smallest_node = node.right
while smallest_node.left:
parent_node = smallest_node
smallest_node = smallest_node.left

# The right Node is the smallest one
if parent_node == node:
smallest_node.left = node.left

# The smallest Node was found to the left of its right branch
else:
parent_node.left = smallest_node.right
smallest_node.left = node.left
smallest_node.right = node.right
return smallest_node

elif node.value > value and node.left:
node.left = self.__remove(node.left, value)

elif node.value < value and node.right:
node.right = self.__remove(node.right, value)

node.compute_size()
return node

def __contains(self, node, value):
if node.value == value:
return True
if node.value > value and node.left:
return self.__contains(node.left, value)
if node.value < value and node.right:
return self.__contains(node.right, value)
return False

if __name__ == "__main__":
tree = BSTree()

values = [100, 50, 150, 25, 75, 120, 200, 110, 115]

for v in values:
tree.insert(v)

print("####")
tree.pre_order_traversal()

tree.remove(100)

print("####")
tree.pre_order_traversal()
• Are you targeting Python 2, Python 3, or both? If Python 3, do you need compatibility with Python < 3.3? Sep 17 '17 at 3:00
• I forgot to mention that I am targeting Python 3. I don't think that I am going to ever use this piece of code in production. Sep 17 '17 at 3:09
• I would suggest that your "Private API" methods don't actually use self. They would probably be better as methods on the node class, such that BSTree.foo calls self.root.foo(), etc. Sep 17 '17 at 14:22
• @AustinHastings Comments are for seeking clarification to the question, and may be deleted. Please put all suggestions for improving the code in answers, even brief remarks. Sep 17 '17 at 18:40

Bugs

You have a size-tracking bug. Also, the .contains() method crashes on an empty tree.

>>> t = BSTree()
>>> t.size()
0
>>> t.insert(3)
>>> t.size()
1
>>> t.remove(3)
>>> t.size()
1
>>> t.contains(3)
Traceback (most recent call last):
... line 170, in __contains
if node.value == value:
AttributeError: 'NoneType' object has no attribute 'value'

Interface

It would be nice if the return value of the .insert() and .remove() methods indicated whether the operations succeeded or failed.

Furthermore, the size of the tree can either increase by 1 if .insert() succeeds, decrease by 1 if .remove() succeeds, or remain the same if either fails. There is no need for BSTNode to have a .compute_size() method.

I suggest defining .size() as a read-only property, so that users can write t.size instead of t.size().

Use a single underscore prefix to indicate private members. Double underscore is for name mangling.

Be consistent in your docstring wording. PEP 257 says that you should prefer the imperative to the indicative:

• The docstring is a phrase ending in a period. It prescribes the function or method's effect as a command ("Do this", "Return that"), not as a description; e.g. don't write "Returns the pathname ...".

Traversal API

You mis-wrote the preamble for .post_order_traversal():

if fn is None:
def fn(x): return

print(x)

But even when written correctly, the preamble is silly. In Python 3, print is a function, so you could just make it the default value:

def post_order_traversal(self, fn=print):
…

For even more idiomatic Python, avoid passing fn as a callback. Instead, do the traversal as a generator function, yielding each value. The caller would write instead:

for value in tree.post_order_traversal():
print(value)

Private helper methods

Each of your private helper method exists to serve one corresponding public method. For situations like this, just write an inner function, to reduce clutter, and to avoid having to jump around when reading the code.

class BSTree:
…

def pre_order_traversal(self):
def pre_order(node):
if node is None: return
yield node.value
yield from pre_order(node.left)
yield from pre_order(node.right)
yield from pre_order(self._root)

…

Suggested solution

For brevity, I have omitted the docstrings.

I've fixed the crash in .contains(), and rewritten its helper method to be more compact.

In __remove(), cases 1 and 2 can be condensed down to two lines. I've also renamed smallest_node to more clearly indicate that it refers to the next larger node in the tree.

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

class BSTree:
def __init__(self):
self._root = None
self._size = 0

def pre_order_traversal(self):
def pre_order(node):
if node is None: return
yield node.value
yield from pre_order(node.left)
yield from pre_order(node.right)
yield from pre_order(self._root)

def in_order_traversal(self):
def in_order(node):
if node is None: return
yield from in_order(node.left)
yield node.value
yield from in_order(node.right)
yield from in_order(self._root)

def post_order_traversal(self):
def post_order(node):
if node is None: return
yield from post_order(node.left)
yield from post_order(node.right)
yield node.value
yield from post_order(self._root)

@property
def size(self):
return self._size

def contains(self, value):
def _contains(node, value):
return (
False if node is None else
_contains(node.left, value) if value < node.value else
_contains(node.right, value) if value > node.value else
True
)
return _contains(self._root, value)

def insert(self, value):
def _insert(node, value):
if node is None:
return BSTNode(value)
elif value == node.value:
return None
elif value < node.value:
node.left = _insert(node.left, value)
elif value > node.value:
node.right = _insert(node.right, value)
return node
self._root = _insert(self._root, value)
if self._root:
self._size += 1
return self._root is not None

def remove(self, value):
def _remove(node, value):
if node.value == value:
if not (node.left and node.right):
return node.left or node.right, True
else:
# Replace the node with its next larger successor
successor, parent = node.right, node
while successor.left:
successor, parent = successor.left, successor

successor.left = node.left
if parent != node:
parent.left = successor.right
successor.right = node.right
return successor, True
elif value < node.value and node.left:
node.left, removed = _remove(node.left, value)
return node, removed
elif value > node.value and node.right:
node.right, removed = _remove(node.right, value)
return node, removed
return node, False
if self._root is None:
return False
self._root, removed = _remove(self._root, value)
self._size -= int(removed)
return removed
• Thank you for taking the time to give me such a thoughtful review. I really liked the idea of returning tuples to indicate an operation's result and using generators instead of callback functions. Sep 17 '17 at 15:20
self.size = self.size + self.left.size

This is more naturally expressed as self.size += self.left.size

self.__root = None
self.__size = 0

The double underscore invokes name mangling, which can be helpful it you anticipate inheritance. Here, a single underscore would probably suffice.

You organized the source code into Public and Private API sections, which clearly makes some sort of sense. Personally I would have found it more helpful, when reading code, to see each public method immediately followed by its private method. For example, your (excellent) docstrings matter a great deal when trying to understand implementations, but one must scroll to find them. Methods like pre_order_traversal() are already using nested functions. Private helpers that start with _ are perfectly nice, but you might consider using the alternative of hiding each helper by making it a nested function.

values = [100, 50, 150, 25, 75, 120, 200, 110, 115]

Thank you for including test code that exercises the target code. That is certainly helpful. It would be more helpful to include a unittest.TestCase with good measured code coverage of the target code.

Overall, it looks solid. Ship it!