Add method added. Also adds convenience method has_child in node class

01_binary_search_trees/bst.py

@@ -15,6 +15,31 @@ class BinarySearchTree():
         self.size = 0
         self.root_node = None
 
+    def add(self, value):
+        if self.size == 0:
+            self.root_node = BinarySearchNode(value, None) # root node has no parent
+
+        else:
+            node = self.root_node
+            is_left = None
+            while(node != None):
+                ptr_node = node # keep parent in memory
+                if node.value > value:  # we go to the left child
+                    node = node.left_child
+                    is_left = True
+                else: # we move to the right child
+                    node = node.right_child
+                    is_left = False
+
+            # we achieved a leaf
+            node = BinarySearchNode(value, ptr_node)
+            if is_left:
+                ptr_node.left_child = node
+            else:
+                ptr_node.right_child = node
+
+        self.size += 1
+
 class BinarySearchNode():
     """
     Defines any node of the Binary Search Tree.
@@ -22,8 +47,15 @@ class BinarySearchNode():
     It also must have a parent
     """
     def __init__(self, value, parent):
-        self.parent_node = parent
-        self.node_value = value
+        self.parent = parent
+        self.value = value
 
         self.left_child = None
         self.right_child = None
+
+    def has_left_child(self):
+        return not(self.left_child is None)
+
+    def has_right_child(self):
+        return self.right_child is not None
+

01_binary_search_trees/bst_test.py

@@ -16,6 +16,49 @@ class test_binary_search_tree(unittest.TestCase):
         bst = BinarySearchTree()
         self.assertEqual(bst.size, 0)
 
+        node_val1 = 6
+        bst.add(node_val1)
+
+        self.assertEqual(bst.size, 1)
+
+        node_val2= 4
+        bst.add(node_val2)
+        node_val3= 10
+        bst.add(node_val3)
+        node_val4= 16
+        bst.add(node_val4)
+        self.assertEqual(bst.size, 4)
+
+    def test_has_children(self):
+
+        #    6
+        #   / \
+        #  4   10
+        #       \
+        #        16
+
+        bst = BinarySearchTree()
+        node_val1 = 6
+        bst.add(node_val1)
+        node_val2= 4
+        bst.add(node_val2)
+        node_val3= 10
+        bst.add(node_val3)
+        node_val4= 16
+        bst.add(node_val4)
+        self.assertEqual(bst.size, 4)
+
+        node = bst.root_node  # we have 6
+        self.assertEqual(node.value, node_val1)
+        print node.left_child, node.value
+        self.assertTrue(node.has_left_child())
+        self.assertEqual(node.left_child.value, node_val2)
+        self.assertTrue(node.has_right_child())
+        self.assertEqual(node.right_child.value, node_val3)
+
+        node = node.left_child # we have 4 now
+        self.assertFalse(node.has_left_child())
+        self.assertFalse(node.has_right_child())
 
 if __name__ == '__main__':
     unittest.main()