Modify delete function and add docstrings to binary_tree

pydsa/binary_tree.py

@@ -1,8 +1,12 @@
-# Binary Tree
-# A binary tree is a hierarchial data structure in which each node has at most
-# two children, which are referred to as the left child and the right child.
-# Insertion and deletion are O(1)
-# Tree traversals are O(n)
+"""
+Binary Tree
+A binary tree is a hierarchial data structure in which each node has at most
+two children, which are referred to as the left child and the right child.
+Insertion is O(1) and deletion is O(n).
+Tree traversals are O(n).
+Reference used: http://geeksquiz.com/binary-search-tree-set-2-delete/
+and http://www.geeksforgeeks.org/write-a-c-program-to-delete-a-tree/
+"""
 
 
 class BTNode(object):
@@ -20,11 +24,15 @@ class BTNode(object):
     [1, 2, 3]
     >>> bt.postorder(bt)
     [2, 3, 1]
-    >>> bt.delete("left")
+    >>> bt.delete(bt.left)
     >>> bt.inorder(bt)
     [1, 3]
     """
     def __init__(self, key, left=None, right=None):
+        """
+        Initializes a node with value key and optionally with left and/or
+        right child nodes.
+        """
         self.left = left
         self.right = right
         self.key = key
@@ -33,23 +41,70 @@ def __init__(self, key, left=None, right=None):
         self.postlist = []
 
     def insert(self, child, key):
+        """
+        Takes in a string child to insert the new node with value key at left
+        or right of current node.
+        """
         childNode = BTNode(key)
         if child == "left":
             self.left = childNode
         elif child == "right":
             self.right = childNode
 
-    def delete(self, child):
-        if child == "left":
-            self.left = None
-        elif child == "right":
-            self.right = None
+    def delete(self, root):
+        self.deleteUtil(self, root)
+
+    def deleteUtil(self, node, root):
+        """
+        Recursively searches sub-trees to find root (the node to be deleted).
+        When found, replaces it with the child node if other is None or the
+        inorder successor (and recursively deletes it) if both child nodes
+        are not None.
+        """
+        if node is None:
+            return node
+
+        node.left = self.deleteUtil(node.left, root)
+        node.right = self.deleteUtil(node.right, root)
+
+        if node == root:
+            if root.left is None:
+                temp = root.right
+                root = None
+                return temp
+
+            elif root.right is None:
+                temp = root.left
+                root = None
+                return temp
+
+            # Get inorder successor of root
+            temp = self.getLeftmost(root.right)
+            root.key = temp.key
+
+            # Recursively delete inorder successor
+            root.right = self.deleteUtil(root.right, temp)
+
+        return node
+
+    def getLeftmost(self, root):
+        """
+        Returns the leftmost node in the tree rooted at root.
+        """
+        current = root
+        while current.left is not None:
+            current = current.left
+        return current
 
     def inorder(self, root):
         self.inlist = []
         return self.inorderUtil(root)
 
     def inorderUtil(self, root):
+        """
+        Recursively traverses left sub-tree, then current node and then the
+        right sub-tree.
+        """
         if root:
             self.inorderUtil(root.left)
             self.inlist.append(root.key)
@@ -61,6 +116,10 @@ def preorder(self, root):
         return self.preorderUtil(root)
 
     def preorderUtil(self, root):
+        """
+        Traverses the current node, then recursively traverses left sub-tree,
+        and then the right sub-tree.
+        """
         if root:
             self.prelist.append(root.key)
             self.preorderUtil(root.left)
@@ -72,6 +131,10 @@ def postorder(self, root):
         return self.postorderUtil(root)
 
     def postorderUtil(self, root):
+        """
+        Recursively traverses left sub-tree, then the right sub-tree and then
+        the current node.
+        """
         if root:
             self.postorderUtil(root.left)
             self.postorderUtil(root.right)

pydsa/tests/test_binary_tree.py

@@ -5,28 +5,28 @@ def test_binary_tree():
     bt = BTNode(1)
     bt.insert("left", 2)
     bt.insert("right", 3)
+    bt.right.insert("left", 6)
+    bt.right.insert("right", 7)
     bt.left.insert("left", 4)
     bt.left.insert("right", 5)
-    bt.left.right.insert("left", 6)
-    bt.left.right.insert("right", 7)
-    bt.right.insert("right", 8)
-    bt.right.right.insert("left", 9)
+    bt.left.left.insert("left", 8)
+    bt.left.left.insert("right", 9)
 
     inlist = bt.inorder(bt)
     prelist = bt.preorder(bt)
     postlist = bt.postorder(bt)
 
-    assert inlist == [4, 2, 6, 5, 7, 1, 3, 9, 8]
-    assert prelist == [1, 2, 4, 5, 6, 7, 3, 8, 9]
-    assert postlist == [4, 6, 7, 5, 2, 9, 8, 3, 1]
+    assert inlist == [8, 4, 9, 2, 5, 1, 6, 3, 7]
+    assert prelist == [1, 2, 4, 8, 9, 5, 3, 6, 7]
+    assert postlist == [8, 9, 4, 5, 2, 6, 7, 3, 1]
 
-    bt.left.right.delete("left")
-    bt.right.delete("right")
+    bt.delete(bt.left)
+    bt.delete(bt.right)
 
     inlist = bt.inorder(bt)
     prelist = bt.preorder(bt)
     postlist = bt.postorder(bt)
 
-    assert inlist == [4, 2, 5, 7, 1, 3]
-    assert prelist == [1, 2, 4, 5, 7, 3]
-    assert postlist == [4, 7, 5, 2, 3, 1]
+    assert inlist == [8, 4, 9, 5, 1, 6, 7]
+    assert prelist == [1, 5, 4, 8, 9, 7, 6]
+    assert postlist == [8, 9, 4, 5, 6, 7, 1]