From 7215f01258eebd5635b0d49b278f1cee2b598d6d Mon Sep 17 00:00:00 2001 From: Yash Kumar Saini <115717039+yashksaini-coder@users.noreply.github.com> Date: Thu, 24 Oct 2024 18:41:58 +0000 Subject: [PATCH] Enhanced the BST implementation --- .../BinarySearchTree/bst.py | 181 +++++++++++++----- .../BinarySearchTree/bstnode.py | 12 +- .../BinarySearchTree/main.py | 98 ++++++---- .../BinarySearchTree/tree_exceptions.py | 7 + 4 files changed, 210 insertions(+), 88 deletions(-) create mode 100644 Algorithms_and_Data_Structures/BinarySearchTree/tree_exceptions.py diff --git a/Algorithms_and_Data_Structures/BinarySearchTree/bst.py b/Algorithms_and_Data_Structures/BinarySearchTree/bst.py index c671687253..6832a0c278 100644 --- a/Algorithms_and_Data_Structures/BinarySearchTree/bst.py +++ b/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) - \ No newline at end of file + 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) diff --git a/Algorithms_and_Data_Structures/BinarySearchTree/bstnode.py b/Algorithms_and_Data_Structures/BinarySearchTree/bstnode.py index aeef4c901a..1f305fe486 100644 --- a/Algorithms_and_Data_Structures/BinarySearchTree/bstnode.py +++ b/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 \ No newline at end of file + 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})" diff --git a/Algorithms_and_Data_Structures/BinarySearchTree/main.py b/Algorithms_and_Data_Structures/BinarySearchTree/main.py index 927bf53ed2..5e6b10c3af 100644 --- a/Algorithms_and_Data_Structures/BinarySearchTree/main.py +++ b/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() \ No newline at end of file diff --git a/Algorithms_and_Data_Structures/BinarySearchTree/tree_exceptions.py b/Algorithms_and_Data_Structures/BinarySearchTree/tree_exceptions.py new file mode 100644 index 0000000000..bd24bb21e6 --- /dev/null +++ b/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