Enhanced the BST implementation

Algorithms_and_Data_Structures/BinarySearchTree/bst.py

@@ -1,58 +1,143 @@
-from bstnode import BSTnode
+from typing import List, Optional, Tuple
+from bstnode import BSTNode
+from tree_exceptions import EmptyTreeError, InvalidValueError
 
 class BST:
     def __init__(self):
         self.root = None
         self.num_nodes = 0
-    
-    def insert(self ,value:int) -> None:
-        if self.root is None:
-            newnode = BSTnode(value)
-            self.root = newnode
-        else:
-           self.insert_helper(self.root ,value)
+
+    def insert(self, value: int) -> None:
+        """Insert a value into the BST, handling duplicates by incrementing count"""
+        if not isinstance(value, (int, float)):
+            raise InvalidValueError("Value must be a number")
+
+        self.root = self._insert_helper(self.root, value)
         self.num_nodes += 1
-
-    def insert_helper(self ,root:BSTnode ,value:int) -> None:
-        if value < root.value:
-            if root.leftPtr is None:
-                root.leftPtr = BSTnode(value)
-            else:
-                self.insert_helper(root.leftPtr, value)
+
+    def _insert_helper(self, root: Optional[BSTNode], value: int) -> BSTNode:
+        if root is None:
+            return BSTNode(value)
+
+        if value == root.value:
+            root.count += 1
+        elif value < root.value:
+            root.leftPtr = self._insert_helper(root.leftPtr, value)
         else:
-            if root.rightPtr is None:
-                root.rightPtr = BSTnode(value)
-            else:
-                self.insert_helper(root.rightPtr, value)
+            root.rightPtr = self._insert_helper(root.rightPtr, value)
 
-    def search(self ,value:int) -> bool:
-        temp_node = self.root
-        while temp_node is not None:
-            if temp_node.value == value:
-                return True
-            elif temp_node.value > value:
-                temp_node = temp_node.leftPtr
-            else:
-                temp_node = temp_node.rightPtr
-        return False
-
-    def inorder(self) -> None:
-        self.inorder_helper(self.root)
-
-    def inorder_helper(self ,root:BSTnode) -> None:
-        if root is not None:
-            self.inorder_helper(root.leftPtr)
-            print(root.value, end=' ')
-            self.inorder_helper(root.rightPtr)
-
-    def get_height(self):
-        return self.height_helper(self.root)
-
-    def height_helper(self, root: BSTnode) -> int:
+        root.height = 1 + max(self._get_height(root.leftPtr), 
+                            self._get_height(root.rightPtr))
+        return root
+
+    def delete(self, value: int) -> None:
+        """Delete a value from the BST"""
+        if self.root is None:
+            raise EmptyTreeError("Cannot delete from empty tree")
+        self.root = self._delete_helper(self.root, value)
+        self.num_nodes -= 1
+
+    def _delete_helper(self, root: Optional[BSTNode], value: int) -> Optional[BSTNode]:
         if root is None:
-            return -1 
+            return None
+
+        if value < root.value:
+            root.leftPtr = self._delete_helper(root.leftPtr, value)
+        elif value > root.value:
+            root.rightPtr = self._delete_helper(root.rightPtr, value)
         else:
-            left_height = self.height_helper(root.leftPtr)
-            right_height = self.height_helper(root.rightPtr)
-            return 1 + max(left_height, right_height)
-
+            if root.count > 1:
+                root.count -= 1
+                return root
+
+            if root.leftPtr is None:
+                return root.rightPtr
+            elif root.rightPtr is None:
+                return root.leftPtr
+
+            # Node with two children
+            successor = self._find_min(root.rightPtr)
+            root.value = successor.value
+            root.count = successor.count
+            successor.count = 1
+            root.rightPtr = self._delete_helper(root.rightPtr, successor.value)
+
+        root.height = 1 + max(self._get_height(root.leftPtr), 
+                            self._get_height(root.rightPtr))
+        return root
+
+    def search(self, value: int) -> Tuple[bool, int]:
+        """Search for a value and return (found, count)"""
+        node = self._search_helper(self.root, value)
+        return (True, node.count) if node else (False, 0)
+
+    def _search_helper(self, root: Optional[BSTNode], value: int) -> Optional[BSTNode]:
+        if root is None or root.value == value:
+            return root
+
+        if value < root.value:
+            return self._search_helper(root.leftPtr, value)
+        return self._search_helper(root.rightPtr, value)
+
+    def get_height(self) -> int:
+        """Get the height of the tree"""
+        return self._get_height(self.root)
+
+    def _get_height(self, node: Optional[BSTNode]) -> int:
+        return node.height if node else 0
+
+    def _find_min(self, node: BSTNode) -> BSTNode:
+        current = node
+        while current.leftPtr:
+            current = current.leftPtr
+        return current
+
+    def get_balance(self, node: Optional[BSTNode]) -> int:
+        """Get balance factor of a node"""
+        if node is None:
+            return 0
+        return self._get_height(node.leftPtr) - self._get_height(node.rightPtr)
+
+    def traversals(self) -> dict:
+        """Return all three traversals in a dictionary"""
+        return {
+            'inorder': self.inorder(),
+            'preorder': self.preorder(),
+            'postorder': self.postorder()
+        }
+
+    def inorder(self) -> List[int]:
+        """Return inorder traversal as a list"""
+        result = []
+        self._inorder_helper(self.root, result)
+        return result
+
+    def _inorder_helper(self, root: Optional[BSTNode], result: List[int]) -> None:
+        if root:
+            self._inorder_helper(root.leftPtr, result)
+            result.extend([root.value] * root.count)
+            self._inorder_helper(root.rightPtr, result)
+
+    def preorder(self) -> List[int]:
+        """Return preorder traversal as a list"""
+        result = []
+        self._preorder_helper(self.root, result)
+        return result
+
+    def _preorder_helper(self, root: Optional[BSTNode], result: List[int]) -> None:
+        if root:
+            result.extend([root.value] * root.count)
+            self._preorder_helper(root.leftPtr, result)
+            self._preorder_helper(root.rightPtr, result)
+
+    def postorder(self) -> List[int]:
+        """Return postorder traversal as a list"""
+        result = []
+        self._postorder_helper(self.root, result)
+        return result
+
+    def _postorder_helper(self, root: Optional[BSTNode], result: List[int]) -> None:
+        if root:
+            self._postorder_helper(root.leftPtr, result)
+            self._postorder_helper(root.rightPtr, result)
+            result.extend([root.value] * root.count)

Algorithms_and_Data_Structures/BinarySearchTree/bstnode.py

@@ -1,6 +1,10 @@
-class BSTnode:
-
-    def __init__(self ,value:int) -> None:
+class BSTNode:
+    def __init__(self, value: int) -> None:
         self.value = value
         self.leftPtr = None
-        self.rightPtr = None
+        self.rightPtr = None
+        self.height = 1  # simple balance factor calculation
+        self.count = 1   # For handling duplicate values
+
+    def __str__(self) -> str:
+        return f"Node(value={self.value}, height={self.height}, count={self.count})"

Algorithms_and_Data_Structures/BinarySearchTree/main.py

@@ -1,53 +1,79 @@
 from bst import BST
+from tree_exceptions import EmptyTreeError, InvalidValueError
+
+def print_menu():
+    print("\n" + "="*40)
+    print("Enhanced Binary Search Tree Menu")
+    print("="*40)
+    print("1. Insert a value")
+    print("2. Delete a value")
+    print("3. Search for a value")
+    print("4. Display all traversals")
+    print("5. Get tree height")
+    print("6. Get total nodes")
+    print("7. Exit")
+    print("="*40)
 
 def main():
     bst = BST()
 
     while True:
-        print("\n" + "="*30)
-        print("Binary Search Tree Menu")
-        print("="*30)
-        print("1. Insert a value")
-        print("2. Search for a value")
-        print("3. Display in-order traversal")
-        print("4. Get tree height")
-        print("5. Exit")
-        print("="*30)
-
+        print_menu()
         try:
-            choice = int(input("Enter your choice (1-5): "))
+            choice = int(input("Enter your choice (1-7): "))
         except ValueError:
-            print("Invalid input. Please enter a number between 1 and 5.")
+            print("Invalid input. Please enter a number between 1 and 7.")
             continue
 
-        if choice == 1:
-            value = int(input("Enter the value to insert: "))
-            bst.insert(value)
-            print(f"Value {value} inserted into the BST.")
+        try:
+            if choice == 1:
+                value = int(input("Enter the value to insert: "))
+                bst.insert(value)
+                print(f"Value {value} inserted into the BST.")
 
-        elif choice == 2:
-            value = int(input("Enter the value to search: "))
-            found = bst.search(value)
-            if found:
-                print(f"Value {value} found in the BST.")
-            else:
-                print(f"Value {value} not found in the BST.")
+            elif choice == 2:
+                if bst.root is None:
+                    print("Tree is empty!")
+                    continue
+                value = int(input("Enter the value to delete: "))
+                bst.delete(value)
+                print(f"Value {value} deleted from the BST.")
 
-        elif choice == 3:
-            print("In-order traversal of the BST: ", end="")
-            bst.inorder()
-            print()  
+            elif choice == 3:
+                value = int(input("Enter the value to search: "))
+                found, count = bst.search(value)
+                if found:
+                    print(f"Value {value} found in the BST. Count: {count}")
+                else:
+                    print(f"Value {value} not found in the BST.")
 
-        elif choice == 4:
-            height = bst.get_height()
-            print(f"Height of the BST: {height}")
+            elif choice == 4:
+                if bst.root is None:
+                    print("Tree is empty!")
+                    continue
+                traversals = bst.traversals()
+                print("\nInorder:", traversals['inorder'])
+                print("Preorder:", traversals['preorder'])
+                print("Postorder:", traversals['postorder'])
 
-        elif choice == 5:
-            print("Exiting program.")
-            break
+            elif choice == 5:
+                height = bst.get_height()
+                print(f"Height of the BST: {height}")
+
+            elif choice == 6:
+                print(f"Total nodes in the BST: {bst.num_nodes}")
+
+            elif choice == 7:
+                print("Thank you for using the BST program!")
+                break
+
+            else:
+                print("Invalid choice. Please select a number between 1 and 7.")
 
-        else:
-            print("Invalid choice. Please select a number between 1 and 5.")
+        except (EmptyTreeError, InvalidValueError) as e:
+            print(f"Error: {e}")
+        except Exception as e:
+            print(f"An unexpected error occurred: {e}")
 
 if __name__ == "__main__":
-    main()
+    main()

Algorithms_and_Data_Structures/BinarySearchTree/tree_exceptions.py

@@ -0,0 +1,7 @@
+class EmptyTreeError(Exception):
+    """Raised when operation is performed on empty tree"""
+    pass
+
+class InvalidValueError(Exception):
+    """Raised when invalid value is provided"""
+    pass