This notebook was prepared by Donne Martin. Source and license info is on GitHub.
Insert will be tested through the following traversal:
If the root
input is None
, return a tree with the only element being the new root node.
You do not have to code the in-order traversal, it is part of the unit test.
Complexity:
%%writefile bst.py
class Node(object):
def __init__(self, data):
self.data = data
self.left = None
self.right = None
self.parent = None
def __repr__(self):
return str(self.data)
class Bst(object):
def __init__(self, root=None):
self.root = root
def insert(self, data):
if data is None:
raise TypeError('data cannot be None')
if self.root is None:
self.root = Node(data)
return self.root
else:
return self._insert(self.root, data)
def _insert(self, node, data):
if node is None:
return Node(data)
if data <= node.data:
if node.left is None:
node.left = self._insert(node.left, data)
node.left.parent = node
return node.left
else:
return self._insert(node.left, data)
else:
if node.right is None:
node.right = self._insert(node.right, data)
node.right.parent = node
return node.right
else:
return self._insert(node.right, data)
Overwriting bst.py
%run bst.py
%run dfs.py
%run ../utils/results.py
%%writefile test_bst.py
import unittest
class TestTree(unittest.TestCase):
def __init__(self, *args, **kwargs):
super(TestTree, self).__init__()
self.results = Results()
def test_tree_one(self):
bst = Bst()
bst.insert(5)
bst.insert(2)
bst.insert(8)
bst.insert(1)
bst.insert(3)
in_order_traversal(bst.root, self.results.add_result)
self.assertEqual(str(self.results), '[1, 2, 3, 5, 8]')
self.results.clear_results()
def test_tree_two(self):
bst = Bst()
bst.insert(1)
bst.insert(2)
bst.insert(3)
bst.insert(4)
bst.insert(5)
in_order_traversal(bst.root, self.results.add_result)
self.assertEqual(str(self.results), '[1, 2, 3, 4, 5]')
print('Success: test_tree')
def main():
test = TestTree()
test.test_tree_one()
test.test_tree_two()
if __name__ == '__main__':
main()
Overwriting test_bst.py
%run -i test_bst.py
Success: test_tree