add methods to the binary tree class for convenience while at the same time keeping the flexibility to call those methods outside by passing the root node.

laurentluce · laurentluce · commit 8522057d20e8 · 2010-11-24T11:50:13.000-08:00
diff --git a/binary_tree.py b/binary_tree.py
@@ -25,6 +25,26 @@ def __init__(self):
   def get_root(self):
     return self.root
 
+  def insert(self, data):
+    node = bt_insert(self.root, data)
+    if self.root == None:
+      self.root = node
+
+  def lookup(self, data):
+    return bt_lookup(self.root, data)
+
+  def delete(self, data):
+    return bt_delete(self.root, data)
+
+  def compare_trees(self, root):
+    return bt_compare_trees(self.root, root)
+
+  def print_tree(self):
+    bt_print_tree(self.root)
+
+  def tree_data(self):
+    return bt_tree_data(self.root)
+
 def bt_insert(root, data):
   """
   Insert new node with data
@@ -59,44 +79,43 @@ def bt_lookup(root, data, parent=None):
   else:
     return root, parent
 
-def bt_compare_trees(t1, t2):
+def bt_compare_trees(root1, root2):
   """
   Compare 2 trees
 
-  @param t1 tree 1
-  @param t2 tree 2
+  @param root1 tree 1 root node
+  @param root2 tree 2 root node
   @returns True if the 2 trees are identical or False
   """
-  if t1 == None and t2 == None:
+  if root1 == None and root2 == None:
     return True
-  elif (t1 == None and t2 != None) or (t1 != None and t2 == None):
+  elif (root1 == None and root2 != None) or (root1 != None and root2 == None):
     return False
   else:
-    return (t1.data == t2.data) \
-        and bt_compare_trees(t1.left, t2.left) \
-        and bt_compare_trees(t1.right, t2.right)
+    return (root1.data == root2.data) \
+        and bt_compare_trees(root1.left, root2.left) \
+        and bt_compare_trees(root1.right, root2.right)
 
 def bt_print_tree(root):
   """
-  Print tree content in order
+  Print tree content inorder
 
   @param root tree root node
   """
-  if root == None:
-    pass
-  else:
+  if root:
     bt_print_tree(root.left)
     print root.data,
     bt_print_tree(root.right)
 
-def bt_tree_data(node):
+def bt_tree_data(root):
   """
   Generator to get the tree nodes data
 
   @param node tree root node
   """
   # we use a stack to traverse the tree in a non-recursive way
   stack = []
+  node = root
   while stack or node: 
     if node:
       stack.append(node)
diff --git a/tests/test_binary_tree.py b/tests/test_binary_tree.py
@@ -8,60 +8,58 @@ def test_binary_tree(self):
     data = [10, 5, 15, 4, 7, 13, 17, 11, 14]
     # create 2 trees with the same content
     bt = binary_tree.BinaryTree()
-    root = binary_tree.bt_insert(bt.get_root(), data[0])
-    for i in data[1:]:
-      binary_tree.bt_insert(root, i)
+    for i in data:
+      bt.insert(i)
 
     bt2 = binary_tree.BinaryTree()
-    root2 = binary_tree.bt_insert(bt2.get_root(), data[0])
-    for i in data[1:]:
-      binary_tree.bt_insert(root2, i)
+    for i in data:
+      bt2.insert(i)
 
     # check if both trees are identical
-    self.assertTrue(binary_tree.bt_compare_trees(root, root2))
+    self.assertTrue(bt.compare_trees(bt2.get_root()))
 
     # check the content of the tree inorder
     t = []
-    for d in binary_tree.bt_tree_data(root):
+    for d in bt.tree_data():
       t.append(d)
     self.assertEquals(t, [4, 5, 7, 10, 11, 13, 14, 15, 17])
 
     # test lookup
-    node, parent = binary_tree.bt_lookup(root, 9)
+    node, parent = bt.lookup(9)
     self.assertTrue(node == None)
-    node, parent = binary_tree.bt_lookup(root, 11)
+    node, parent = bt.lookup(11)
     self.assertTrue(node.data == 11)
     self.assertTrue(parent.data == 13)
 
     # delete a leaf node
-    binary_tree.bt_delete(root, 4)
+    bt.delete(4)
     # check the content of the tree inorder
     t = []
-    for d in binary_tree.bt_tree_data(root):
+    for d in bt.tree_data():
       t.append(d)
     self.assertEquals(t, [5, 7, 10, 11, 13, 14, 15, 17])
 
     # delete a node with 1 child
-    binary_tree.bt_delete(root, 5)
+    bt.delete(5)
     # check the content of the tree inorder
     t = []
-    for d in binary_tree.bt_tree_data(root):
+    for d in bt.tree_data():
       t.append(d)
     self.assertEquals(t, [7, 10, 11, 13, 14, 15, 17])
 
     # delete a node with 2 childs
-    binary_tree.bt_delete(root, 13)
+    bt.delete(13)
     # check the content of the tree inorder
     t = []
-    for d in binary_tree.bt_tree_data(root):
+    for d in bt.tree_data():
       t.append(d)
     self.assertEquals(t, [7, 10, 11, 14, 15, 17])
 
     # delete a node with 2 childs
-    binary_tree.bt_delete(root, 15)
+    bt.delete(15)
     # check the content of the tree inorder
     t = []
-    for d in binary_tree.bt_tree_data(root):
+    for d in bt.tree_data():
       t.append(d)
     self.assertEquals(t, [7, 10, 11, 14, 17])
 
