# Implement BinarySearchTree class in python

I am currently study the basic data structure and trying to implement everything as I go. can anyone give me some feedback about class BinarySearchTree how to make the code more elegant. any review are appreciated.

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

class BinarySearchTree:
"""
Database strucutre for binary search tree
1:search
2:insert
3:delete
4:get_height
5:get_min_value
6:get_max_value
"""

def __init__(self, root=None):
self.root = root

def __iter(self, cur):
if cur is not None:
yield from self.__iter(cur.left)
yield cur
yield from self.__iter(cur.right)

def __repr__(self):
cur = self.root
return ''.join(str(i.val) for i in self.__iter(cur))

def insert(self, key):
"""insert node into binary tree based on node's value"""
cur = self.root
if cur is None:
self.root = TreeNode(key)
return
while cur is not None:
if key < cur.val:
if cur.left is None:
cur.left = TreeNode(key)
return
else:
cur = cur.left
elif key > cur.val:
if cur.right is None:
cur.right = TreeNode(key)
return
else:
cur = cur.right

def search(self, key):
"""find node from binary tree based on node's value"""
cur = self.root
if cur is None:
while cur is not None:
if key < cur.val:
cur = cur.left
elif key > cur.val:
cur = cur.right
else:
return cur

def get_min_value(self):
"""return the min value from tree"""
cur = self.root
while cur is not None and cur.left is not None:
cur = cur.left
return cur.val

def get_max_value(self):
"""return the max value from tree"""
cur = self.root
while cur is not None and cur.right is not None:
cur = cur.right
return cur.val

def get_height(self):
"""return tree height of binary search tree"""
h = 0
return self.__get_height(self.root, h) if self.root else h

def __get_height(self, cur, h):
"""recursion the tree with left subtree and right subtree"""
if cur is None: return h
left_height = self.__get_height(cur.left, h + 1)
right_height = self.__get_height(cur.right, h + 1)
return max(left_height, right_height)

def delete(self, key):
"""delete the node from binary tree based on node's key value"""
self.__delete(self.root, key)

def __delete(self, cur, key):
"""recursion the tree to find the node and delete from tree"""
if cur is None: return
if key < cur.val:
cur.left = self.__delete(cur.left, key)
elif key > cur.val:
cur.right = self.__delete(cur.right, key)
else:
if cur.left is None:
return cur.right
elif cur.right is None:
return cur.left
else:
def __get_successor(n):
while n is not None and n.left is not None:
n = n.left
return n

successor = __get_successor(cur)
cur.key = successor.key
cur.right = self.__delete(cur.right, successor.key)
return cur

if __name__ == '__main__':
bst = BinarySearchTree()
bst.insert(6)
bst.insert(2)
bst.insert(8)
bst.insert(0)
bst.insert(4)
bst.insert(7)
bst.insert(9)
bst.insert(3)
bst.insert(5)
print(bst.search(5).val == 5)
print(bst.search(0).val == 0)
print(bst.search(9).val == 9)
print(bst.search(6).val == 6)
try:
bst.search(13)
except KeyError as e:
print(e)
print(bst.get_height() == 4)
bst.delete(5)
print(bst.get_height() == 4)
print(bst.get_max_value() == 9)
print(bst.get_min_value() == 0)
bst.delete(3)
bst.delete(7)
bst.delete(9)
print(bst.get_height() == 3)
print(bst)


• I believe that if you try to insert a duplicate value, your code will go into an infinite loop. – Austin Hastings Apr 11 '19 at 5:34
• yes, for binary search tree. i am not sure how I should handle duplicated value and how to remove the duplicated recursion based on search, insert and delete – A.Lee Apr 11 '19 at 15:00
• @A.Lee This depends on a specification, but I would expect that infinite loop is unlikely to be desirable. Throwing an exception or not inserting duplicates seems to be preferable solutions. – reducing activity Apr 15 '19 at 15:56

The code is clean, but there are some problems: 1. your binary search tree is not a standard binary search tree; 2. your implementation is not consistent; 3. various issues.

# Problems

## What is a binary search tree?

There is one bug that was pointed in a comment: what happens if you insert a value that is already present in the tree? In the current code, you fall in an infinite loop since the case is simply ignored:

# insert
while cur is not None:
if key < cur.val:
...
elif key > cur.val:
...
# What if key == cur.val???
# Nothing: just loop, and loop, and loop...


But the question is: what would you like to do if the value to insert is already present? Just let the tree as it is, or insert the value and get it twice in the tree? In the standard binary tree, nodes have a key and a value, and you insert a pair key/value: if the key is already present, the previous value is replaced by the new value. You can create a simplified version where the key is the value value itself, and thus just ignore the duplicate values, or allow duplicate values, but I think the best is to create a regular BST with keys and values since it's as easy as create one with only values.

## Implementation consistency

You do not choose between a recursive implementation and an iterative one: insert, search, get_min/max are iterative while get_height and delete are recursive. For a first implementation, I would use only recursion because it's easier to understand. Once the code works, you can improve the speed by removing recursion if and only if it is mandatory.

## Misc

### Double underscores

Do not use them as a synonym for private in other OOP languages:

Generally, double leading underscores should be used only to avoid name conflicts with attributes in classes designed to be subclassed ~ PEP8

### __repr__

For many types, this function makes an attempt to return a string that would yield an object with the same value when passed to eval(), otherwise the representation is a string enclosed in angle brackets that contains the name of the type of the object together with additional information often including the name and address of the object. ~ repr

Your implementation matches more __str__ than __repr__.

### __iter

This method might be confused with __iter__. You should etiehr implement __iter__ or choose another name.

# Tests

## Testing framework

You gave us some tests: that's a very good point. But they are non standard. You should use a testing framework. There are two testing tools shipped with Python: unittest and doctest. I'm really in love with doctest, because can test your module without writing boilerplate code. But this method has some limits: if you want to thoroughly test a module, you'll have to write separate unit tests.

## Code refactoring

When you refactor any code, I think it's very important to make sure the exisiting tests still pass. Hence, I insert a docstring at the top the file with your tests:

"""
>>> bst = BinarySearchTree()
>>> for v in [6,2,8,0,4,7,9,3,5]:
...     bst.insert(v)
>>> [bst.search(v).value for v in [5,0,9,6]]
[5, 0, 9, 6]
>>> bst.search(13)
Traceback (most recent call last):
...
>>> bst.get_height()
4
>>> bst.delete(5)
>>> bst.get_height()
4
>>> bst.get_max_value()
9
>>> bst.get_min_value()
0
>>> for v in [3,7,9]:
...     bst.delete(v)
>>> bst.get_height()
3
>>> bst
02468
"""

# BODY OF THE MODULE

if __name__ == '__main__':
import doctest
doctest.testmod()


Now we can confidently refactor the code: it will perform at least as good as before.

As you see, the tests mimic a REPL, which is very intuitive and well known by every Python user. Usually, you put comments around the tests to explain whats going on.

I made some little improvements to the tests: * bunch insert of values with a for loop * a list comprehension to test the search method

We'll use doctest in the docstrings of functions too.

# A recursive implementation

As I wrote above, a recursive implementation seems a good start, since the trees are inherently recursive structures.

## The TreeNode class

Your __init__ method allows the value to be None.

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


But there is no reason to accept nodes with a None value. Sooner or later, you'll get something like:

key > cur.value


where cur.value is None! Result:

TypeError: '>' not supported between instances of 'xxx' and 'NoneType'


I like the idea that the user of the class is not a child, and won't use None as a value (we should warn him/her). That's the responsibity of the user not to use None -- the user means also: You, when you use your own classes. But your responsibility is to ensure that the object state is correct with the default parameters. That's not the case.

As written, to get a regular BST, we need to add a key field. Since the TreeNode has no method and is just a convenient way to store the attributes of a node, we can use the new 3.7 dataclass:

from dataclasses import dataclass
from typing import Any

@dataclass
class TreeNode:
key: Any
value: Any
left: 'TreeNode' = None
right: 'TreeNode' = None


Note on Any type: we could enforce the presence of __eq__ and __lt__ in the key, but that would be overkill I think. It's Python, not Java!

## search

This is, in my mind, the best entry point: we want a tree to search key/values pairs. Note that, in your code:

if cur is None:
while cur is not None:
...



you don't need to test if cur is None before the loop since the loop will be skipped immediately and the error raised. But we need this test in the recursive version:

class BinarySearchTree:
...

def search(self, key):
"""Return the value assiociated with key, or raise a KeyError exception"""
return _search(self.root, key).value

# standalone function, out of the BinarySearchTree class
def _search(node, key):
"""Search a node by key. Raise a KeyError exception if the key is not in the tree"""
if node is None: # end of recursion
if key < node.key:
return _search(node.left, key)
elif key > node.key:
return _search(node.right, key)
else: # key == node.key
return node


I think we can solve a little puzzle here. It's not an accident you had to create helper methods (__get_height and __delete) for the recursive methods. Actually, the methods get_height and delete are only bootstraps for those methods. Here, BinarySearchTree.search is a bootstrap for the standalone _search function.

That's because, conceptually, every node is also a tree, and a complete tree. That means that, in your code, the tree is represented by the TreeNode class, not the BinarySearchTree class. The BinarySearchTree class is just a wrapper around the root and a method supplier.

It would be possible to attach the _search method to the TreeNode class, but we would have to test for None value at the children level:

class TreeNode:
...

def search(self, key):
"""Search a node by key. Raise a KeyError exception if the key is not in the tree"""
if key < self.key and self.left is not None:
return node.left.search(key)
elif key > node.key and self.right is not None:
return node.right.search(key)
elif key == node.key:
return node


I would probably write it like that in Java, but its easier to use standalone functions in Python.

Now we have the _search function, we can test it. Just add to the docstring a few lines:

"""Search a node by key. Raise a KeyError exception if the key is not in the tree

>>> _search(None, 1)
Traceback (most recent call last):
...
>>> _search(TreeNode(0, 0), 1)
Traceback (most recent call last):
...
>>> _search(TreeNode(1, 0), 1)
TreeNode(key=1, value=0, left=None, right=None)
>>> _search(TreeNode(0, 0, None, TreeNode(1, 1)), 1)
TreeNode(key=1, value=1, left=None, right=None)
"""


We have to fix the tests in the module docstring too, but I keep it for later.

## insert

We know how to search a node by key. Now, we need to insert key/value pairs in the tree:

class BinarySearchTree:
...

def insert(self, key, value):
"""insert node into binary tree based on node's key"""
self.root = _insert(self.root, key, value)

def _insert(node, key, value):
"""Return node extended with a new key/value pair"""
if node is None:
return TreeNode(key, value)

if key < node.key:
return TreeNode(node.key, node.value, _insert(node.left, key, value), node.right)
elif key > node.key:
return TreeNode(node.key, node.value, node.left, _insert(node.right, key, value))
else: # key == node.key
return TreeNode(node.key, value, node.left, node.right)


There is an interesting pattern here:

return TreeNode(node.key, node.value, _insert(node.left, key, value), node.right)


Seems equivalent to:

node.right = _insert(node.left, key, value)
return node


But there is a big difference: the former is side effect free, while the latter is not. I prefer to avoid side effects because the code is easier to understand, but there is a performance cost. Again, we add some tests:

"""Return node extended with a new key/value pair

>>> node = TreeNode(0,0)
>>> for i in range(2):
...     node = _insert(node, i, 2*i+1)
>>> node
TreeNode(key=0, value=1, left=None, right=TreeNode(key=1, value=3, left=None, right=None))
>>> _insert(node, -1, -2)
TreeNode(key=0, value=1, left=TreeNode(key=-1, value=-2, left=None, right=None), right=TreeNode(key=1, value=3, left=None, right=None))
"""


## get_min, get_max and get_height

You should known how to proceed: get_min and get_max are easy to write recursively.

Your implementation of get_height, is recursive, but uses a tail call optimization (a mechanism that prevents the stack from growing insanely). I don't know if you did it on purpose, but I will remove this optimization for the sake of clarity:

def _get_height(node):
"""return tree height of binary search tree"""
if node is None:
return 0
return 1 + max(_get_height(node.left), _get_height(node.right))


## delete

The delete operation needs the two following steps: 1. find the node having the given key; 2. remove the key. The first step is easy, now but the second one is not. Let's look at your code to understand what happens:

    else: # key == cur.key
if cur.left is None:
return cur.right
elif cur.right is None:
return cur.left
else:
def __get_successor(n):
while n is not None and n.left is not None:
n = n.left
return n

successor = __get_successor(cur)
cur.key = successor.key
cur.right = self.__delete(cur.right, successor.key)


If the node lacks one of its chidren, you just return the other child. That's ok.

But if the node has both left and right children, you take the leftmost element of the right child (you could take the rightmost element of the left child) and replace the current node with that node. That's ok because the leftmost element has a value that is: greater than any value in the left child; lower than any other value of the right child (definition of a BST). Hence, you find the successor and delete it from the right child.

Now in recursive idiom:

else: # key == node.key
if node.left is None:
return node.right
else:
successor, right = _detach_min(node.right)
return TreeNode(successor.key, successor.value, node.left, right)


Wait! what is this _detach_min function? When I write my code, I try to keep the momentum. If I don't know how to write something, I just use a function that does not exist yet. Later, I try to write this function:

def _detach_min(node):
"""Return the min value from the tree and the
rest of the tree"""
if node.left is None:
return node, None

m, r = _detach_min(node.left)
return m, TreeNode(node.key, node.value, r, node.right)


First, we detach the min from the left child, then we return this min and the tree without the min.

## __repr__ and __str__

With the dataclass, __repr__ is almost free:

def __repr__(self):
return f"BinarySearchTree(root={self.root})"


We'll see __str__ and the iteration on key/values later.

# Full code

The code is now complete. I will just adapt the initial tests:

"""
>>> bst = BinarySearchTree()
>>> for v in [6,2,8,0,4,7,9,3,5]:
...     bst.insert(v, v)
>>> [bst.search(v) for v in [5,0,9,6]]
[5, 0, 9, 6]
>>> bst.search(13)
Traceback (most recent call last):
...
>>> bst.get_height()
4
>>> bst.delete(5)
>>> bst.get_height()
4
>>> bst.get_max_value()
9
>>> bst.get_min_value()
0
>>> bst.delete(3)
>>> bst.delete(7)
>>> bst.delete(9)
>>> bst.get_height()
3
>>> bst
BinarySearchTree(root=TreeNode(key=6, value=6, left=TreeNode(key=2, value=2, left=TreeNode(key=0, value=0, left=None, right=None), right=TreeNode(key=4, value=4, left=None, right=None)), right=TreeNode(key=8, value=8, left=None, right=None)))
"""

from dataclasses import dataclass
from typing import Any

@dataclass
class TreeNode:
key: Any
value: Any
left: 'TreeNode' = None
right: 'TreeNode' = None

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

def search(self, key):
"""Return the value assiociated with key, or raise a KeyError exception"""
return _search(self.root, key).value

def insert(self, key, value):
"""insert node into binary tree based on node's key"""
self.root = _insert(self.root, key, value)

def get_min_value(self):
"""return the min value from tree"""
return _get_min(self.root).value

def get_max_value(self):
"""return the min value from tree"""
return _get_max(self.root).value

def get_max_value(self):
"""return the min value from tree"""
return _get_max(self.root).value

def get_height(self):
"""return the height from tree"""
return _get_height(self.root)

def delete(self, key):
"""Delete the node having the given key"""
self.root = _delete(self.root, key)

def __repr__(self):
return f"BinarySearchTree(root={self.root})"

# standalone function, out of the BinarySearchTree class
def _search(node, key):
"""Search a node by key. Raise a KeyError exception if the key is not in the tree

>>> _search(None, 1)
Traceback (most recent call last):
...
>>> _search(TreeNode(0, 0), 1)
Traceback (most recent call last):
...
>>> _search(TreeNode(1, 0), 1)
TreeNode(key=1, value=0, left=None, right=None)
>>> _search(TreeNode(0, 0, None, TreeNode(1, 1)), 1)
TreeNode(key=1, value=1, left=None, right=None)
"""
if node is None: # end of recursion

if key < node.key:
return _search(node.left, key)
elif key > node.key:
return _search(node.right, key)
else:
return node

def _insert(node, key, value):
"""Return node extended with a new key/value pair

>>> node = TreeNode(0,0)
>>> for i in range(2):
...     node = _insert(node, i, 2*i+1)
>>> node
TreeNode(key=0, value=1, left=None, right=TreeNode(key=1, value=3, left=None, right=None))
>>> _insert(node, -1, -2)
TreeNode(key=0, value=1, left=TreeNode(key=-1, value=-2, left=None, right=None), right=TreeNode(key=1, value=3, left=None, right=None))
"""
if node is None:
return TreeNode(key, value)

if key < node.key:
return TreeNode(node.key, node.value, _insert(node.left, key, value), node.right)
elif key > node.key:
return TreeNode(node.key, node.value, node.left, _insert(node.right, key, value))
else: # key == node.key
return TreeNode(node.key, value, node.left, node.right)

def _get_min(node):
"""return the min value from tree
>>> node = TreeNode(0,0)
>>> for i in range(3):
...     node = _insert(node, i, 2*i+1)
>>> _get_min(node).value
1
"""
if node.left is None:
return node

return _get_min(node.left)

def _get_max(node):
"""return the max value from tree
>>> node = TreeNode(0,0)
>>> for i in range(3):
...     node = _insert(node, i, 2*i+1)
>>> _get_max(node).value
5
"""
if node.right is None:
return node

return _get_max(node.right)

def _get_height(node):
"""return tree height of binary search tree
>>> node = TreeNode(0,0)
>>> for i in range(3):
...     node = _insert(node, i, 2*i+1)
>>> _get_height(node)
3
"""
if node is None: # end of the recursion
return 0
return 1 + max(_get_height(node.left), _get_height(node.right))

def _delete(node, key):
"""Return the tree without the node having the given key
>>> node = TreeNode(0,0)
>>> node = _insert(node, 1, 1)
>>> node = _insert(node, -1, 2)
>>> _delete(node, 0)
TreeNode(key=1, value=1, left=TreeNode(key=-1, value=2, left=None, right=None), right=None)
"""
if node is None:
return None
if key < node.key:
return TreeNode(node.key, node.value, _delete(node.left, key), node.right)
elif key > node.key:
return TreeNode(node.key, node.value, node.left, _delete(node.right, key))
else: # key == node.key, end of recursion
if node.left is None:
return node.right
else:
successor, right = _detach_min(node.right)
return TreeNode(successor.key, successor.value, node.left, right)

def _detach_min(node):
"""return the min value from tree
>>> node = TreeNode(0,0)
>>> node = _insert(node, 1, 1)
>>> node = _insert(node, -1, 2)
>>> node = _insert(node, -2, 4)
>>> node = _insert(node, -0.5, 2.5)
>>> _detach_min(node)
(TreeNode(key=-2, value=4, left=None, right=None), TreeNode(key=0, value=0, left=TreeNode(key=-1, value=2, left=None, right=TreeNode(key=-0.5, value=2.5, left=None, right=None)), right=TreeNode(key=1, value=1, left=None, right=None)))
"""
if node.left is None:
return node, None

m, r = _detach_min(node.left)
return m, TreeNode(node.key, node.value, r, node.right)

if __name__ == '__main__':
import doctest
doctest.testmod()


As you can see, the code is not only recursive: it's functional, that means that you don't have any side effect in the functions. Usually, functional code is more readable while imperative code is more efficient. Hence there is no good version: it depends on your needs. What matters is to be consistent.. or not, depending on your needs.

# Conclusion

Here are the more important points to notice:

• A node is a complete tree (all standalone functions deal with nodes as trees).
• The code now decouples the operations on the tree (standalone function) and the operations on the mapping (methods).
• doctest was perhaps used where unitest should have been preferred.

# A step further

I will prove that the second point is important. A binary tree is a mapping. Python provides a frame for mappings: https://docs.python.org/3/library/collections.abc.html#collections.abc.Mapping. Why not try to fit this frame?

We just make BinarySearchTree inherit from Mapping:

from collections.abc import MutableMapping

class BinarySearchTree(MutableMapping):
...


Now, Python requires some methods to be implemented:

TypeError: Can't instantiate abstract class BinarySearchTree with abstract methods __delitem__, __getitem__, __iter__, __len__, __setitem__

• __delitem__ is simply our delete method;
• __getitem__ is our search method;
• __setitem__ is our insert method;
• __iter__ must return an iterator over items;
• __len__ must return the number of nodes of the binary tree;

## __iter__

We'll use a generator to implement the iterator:

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

def _iter(node):
if node is None:
return

yield from _iter(node.left)
yield node.key
yield from _iter(node.right)


## __len__

The method is similar to get_height:

def _len(node):
if node is None:
return 0

return 1 + _len(node.left) + _len(node.right)


## __str__

The __str__ may be implemented to show the sorted dictionary:

def __str__(self):
return "{"+", ".join(f"{k}: {v}" for k, v in self.items())+"}"


## A mapping initializer

We need a convenient way to initialize the binary tree:

def __init__(self, mapping={}):
self.root = None
self.update(mapping)


We could update the __repr__ method to use this initializer.

# Code v2

I omit the parts that where not modified:

"""
>>> bst = BinarySearchTree({v: v for v in [6,2,8,0,4,7,9,3,5]})
>>> [bst[v] for v in [5,0,9,6]]
[5, 0, 9, 6]
>>> bst[13]
Traceback (most recent call last):
...
>>> bst.get_height()
4
>>> del bst[5]
>>> bst.get_height()
4
>>> bst.get_max_value()
9
>>> bst.get_min_value()
0
>>> del bst[3]
>>> del bst[7]
>>> del bst[9]
>>> bst.get_height()
3
>>> str(bst)
'{0: 0, 2: 2, 4: 4, 6: 6, 8: 8}'
"""

...

class BinarySearchTree(MutableMapping):
def __init__(self, mapping={}):
self.root = None
self.update(mapping)

def __getitem__(self, key):
return _search(self.root, key).value

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

def __setitem__(self, key, value):
self.root = _insert(self.root, key, value)

...

def __delitem__(self, key):
self.root = _delete(self.root, key)

def __repr__(self):
return f"BinarySearchTree(root={self.root})"

def __str__(self):
return "{"+", ".join(f"{k}: {v}" for k, v in self.items())+"}"

def __len__(self):
return _len(self.root)

def _iter(node):
if node is None:
return

yield from _iter(node.left)
yield node.key
yield from _iter(node.right)

def _len(node):
if node is None:
return 0

return 1 + _len(node.left) + _len(node.right)

...


Have a look at the tests: the class may be used as any standard Python mapping! The important point is that the interface of the binary tree was adapted to fit with the Mapping interface, but we didn't need to modify the standalone functions. That's a sign that the code was, I hope, correctly decoupled.

• thank you so much. learn so much – A.Lee Apr 18 '19 at 0:15