diff --git a/pydsa/binary_tree.py b/pydsa/binary_tree.py index 751a7ab..c33d842 100644 --- a/pydsa/binary_tree.py +++ b/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) diff --git a/pydsa/tests/test_binary_tree.py b/pydsa/tests/test_binary_tree.py index 99a6213..619fc17 100644 --- a/pydsa/tests/test_binary_tree.py +++ b/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]