Add Binary Tree implementation, tests

pydsa/__init__.py

@@ -12,3 +12,4 @@
 from .counting_sort import counting_sort
 from .stack import Stack
 from .radix_sort import radix_sort
+from .binary_tree import BTNode

pydsa/binary_tree.py

@@ -0,0 +1,142 @@
+"""
+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):
+    """
+    Class to create a Binary Tree Node, insert, delete child nodes and
+    print inorder, preorder and postorder traversals.
+
+    >>> from pydsa import binary_tree
+    >>> bt = binary_tree.BTNode(1)
+    >>> bt.insert("left", 2)
+    >>> bt.insert("right", 3)
+    >>> bt.inorder(bt)
+    [2, 1, 3]
+    >>> bt.preorder(bt)
+    [1, 2, 3]
+    >>> bt.postorder(bt)
+    [2, 3, 1]
+    >>> 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
+        self.inlist = []
+        self.prelist = []
+        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, 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)
+            self.inorderUtil(root.right)
+        return self.inlist
+
+    def preorder(self, root):
+        self.prelist = []
+        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)
+            self.preorderUtil(root.right)
+        return self.prelist
+
+    def postorder(self, root):
+        self.postlist = []
+        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)
+            self.postlist.append(root.key)
+        return self.postlist

pydsa/tests/test_binary_tree.py

@@ -0,0 +1,32 @@
+from pydsa.binary_tree import BTNode
+
+
+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.left.insert("left", 8)
+    bt.left.left.insert("right", 9)
+
+    inlist = bt.inorder(bt)
+    prelist = bt.preorder(bt)
+    postlist = bt.postorder(bt)
+
+    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.delete(bt.left)
+    bt.delete(bt.right)
+
+    inlist = bt.inorder(bt)
+    prelist = bt.preorder(bt)
+    postlist = bt.postorder(bt)
+
+    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]