dsa.tree

  1""" Module containing tree class. """
  2
  3class TreeNode:
  4    """ 
  5    A binary tree node implementation.
  6    """
  7    def __init__(self, value, left=None, right=None):
  8        """ 
  9        TreeNode constructor.
 10        
 11        Args:
 12            value: The value of the node.
 13            left: Reference to the left node.
 14            right: Reference to the right node.
 15        """
 16        self.value = value
 17        self.left = left
 18        self.right = right
 19
 20    def copy(self):
 21        """
 22        Return a copy of the node.
 23        """
 24        new_node = TreeNode(self.value) 
 25        if self.left:
 26            new_node.left = self.left.copy()
 27        if self.right:
 28            new_node.right = self.right.copy()
 29        return new_node
 30    
 31    def print(self, level=0):
 32        """ 
 33        Print the contents of a node.
 34
 35        Args:
 36            level: starting level of node
 37        """
 38        if self.right:
 39            self.right.print(level + 1)
 40        print("   " * level + str(self.value))
 41        if self.left:
 42            self.left.print(level + 1)
 43
 44    def __lt__(self, other):
 45        """
 46        Compare the value of two nodes.
 47        """
 48        return self.value < other.value
 49    
 50    def __repr__(self):
 51        """
 52        Return the string representation of a node.
 53        """
 54        if self.value is None:
 55            return "none"
 56        else:
 57            return str(self.value)
 58
 59class Tree:
 60    """ 
 61    A binary search tree (BST) implementation. Can be treated as a plain binary tree if operations (insert, search, delete) are not used and nodes are set manually.
 62    """
 63    def __init__(self, root=None):
 64        """ 
 65        Args:
 66            root: Root node of the Tree.
 67        """
 68        self.root = root
 69        
 70    def search(self, value):
 71        """ 
 72        Search for a value in the binary search tree.
 73
 74        Args:
 75            value: Value to search for.
 76        
 77        Returns:
 78            node with matching value
 79
 80        Raises:
 81            ValueError: if value is not found
 82        """
 83        current = self.root
 84        
 85        while current is not None:
 86            if value == current.value:
 87                return current
 88            elif value < current.value:
 89                current = current.left
 90            elif value > current.value:
 91                current = current.right
 92
 93        raise ValueError(f"Value {value} not found in the tree")
 94
 95    def insert(self, value):
 96        """ 
 97        Insert a value in the binary search tree.
 98
 99        Args:
100            value: The value to insert.
101        
102        Returns:
103            The root of the tree after deletion.
104
105        Raises:
106            ValueError: if value is already in the tree
107        """
108        self.root = self.insert_rec(self.root, value)
109        return self.root
110    
111    def insert_rec(self, root, value):
112        if root is None:
113            return TreeNode(value)
114        if value < root.value:
115            root.left = self.insert_rec(root.left, value)
116        elif value > root.value:
117            root.right = self.insert_rec(root.right, value)
118        else:
119            raise ValueError("Value already exists in the tree")
120        return root
121
122    def insert_iterative(self, value):
123        """ 
124        Insert a value in the binary search tree (iterative implementation).
125
126        Args:
127            value: The value to insert.
128        
129        Returns:
130            None
131
132        Raises:
133            ValueError: if value is already in the tree
134        """
135        current = self.root
136        if self.root is None:
137            self.root = TreeNode(value)
138            return
139        
140        while current is not None:
141            if value < current.value:
142                if current.left is None:
143                    current.left = TreeNode(value)
144                    return
145                else:
146                    current = current.left
147            elif value > current.value:
148                if current.right is None:
149                    current.right = TreeNode(value)
150                    return
151                else:
152                    current = current.right
153            else:
154                return 
155   
156    def delete(self, value):
157        """ 
158        Delete a value from the binary search tree.
159
160        Args:
161            value: The value to delete.
162        
163        Returns:
164            The root of the tree after insertion.
165
166        
167        Raises:
168            ValueError: if value is not found.
169        """
170        self.root = self.delete_node(self.root, value)
171        return self.root
172        
173    def delete_node(self, node, value):
174        """ 
175        Helper function to delete a value from the binary search tree. (Use delete() instead)
176
177        Args:
178            value: The value to delete.
179            node: The current node.
180
181        Returns:
182            The new subtree root after deletion.
183        
184        Raises:
185            ValueError: if value is not found.
186        """ 
187        if node is None:
188            raise ValueError("Value not found in the tree")
189            return None
190        
191        if value < node.value:
192            node.left = self.delete_node(node.left, value)
193        elif value > node.value:
194            node.right = self.delete_node(node.right, value)
195        else:
196            if node.left is None:
197                subtree = node.right
198                node = None
199                return subtree
200            elif node.right is None:
201                subtree = node.left
202                node = None
203                return subtree
204            
205            successor = self.successor_node(node.right)
206            node.value = successor.value
207            node.right = self.delete_node(node.right, successor.value)
208            
209        return node
210    
211    def successor_node(self, node=None):
212        """ 
213        Return the successor node (the minimum value in a binary search tree's right subtree).
214
215        Args:
216            node: The starting node.
217        
218        Returns:
219            node with the lowest value in the BST
220            None if not found
221
222        Raises:
223            ValueError: if value is not found.
224        """
225        if node is None:
226            node = self.root
227        
228        if node.left is None:
229            return node
230        else:
231            return self.successor_node(node.left)
232    
233    def predecessor_node(self, node=None):
234        """ 
235        Return the predecessor node (the maximum value in a binary search tree's left subtree).
236
237        Args:
238            node: The starting node.
239        
240        Returns:
241            node with the lowest value in the BST
242            None if not found
243
244        Raises:
245            ValueError: if value is not found.
246        """
247        if node is None:
248            node = self.root
249        
250        if node.right is None:
251            return node
252        else:
253            return self.predecessor_node(node.right)
254    
255    def delete_iterative(self, root, value):
256        """ 
257        Delete a value from the binary search tree (iterative version).
258
259        Args:
260            value: The value to delete.
261        
262        Returns:
263            The root of the tree after insertion.
264
265        
266        Raises:
267            ValueError: if value is not found.
268        """
269        parent = None
270        current = root
271
272        # Find the node to delete and its parent
273        while current and current.key != value:
274            parent = current
275            if value < current.key:
276                current = current.left
277            else:
278                current = current.right
279
280        if current is None:
281            return root  # Key not found
282
283        # Case 1 & 2: Node has at most one child
284        def replace_node_in_parent(new_child):
285            if parent is None:
286                return new_child  # Deleting the root node
287            if parent.left == current:
288                parent.left = new_child
289            else:
290                parent.right = new_child
291            return root
292
293        if current.left is None:
294            return replace_node_in_parent(current.right)
295        elif current.right is None:
296            return replace_node_in_parent(current.left)
297
298        # Case 3: Node has two children
299        # Find in-order successor and its parent
300        successor_parent = current
301        successor = current.right
302        while successor.left:
303            successor_parent = successor
304            successor = successor.left
305
306        # Replace current's key with successor's key
307        current.key = successor.key
308
309        # Delete successor node (which has at most one child)
310        if successor_parent.left == successor:
311            successor_parent.left = successor.right
312        else:
313            successor_parent.right = successor.right
314
315        return root
316
317    def print(self):
318        """ 
319        Print the values in the BST.
320        """
321        self.root.print()
322        
323    def __len__(self):
324        """ 
325        Return the number of nodes in the BST.
326        
327        Returns:
328            The number of nodes in the BST.
329        """
330        def count_nodes(node):
331            if node is None:
332                return 0
333            return 1 + count_nodes(node.left) + count_nodes(node.right)
334        
335        return count_nodes(self.root)
336        
337
338    def __eq__(self, other):
339        """
340        Compare two Tree objects for value-based equality (structure and values).
341
342        Returns:
343            True if both are Tree instances and their structures and values are equal, False otherwise.
344        """
345        if not isinstance(other, Tree):
346            return False
347        def nodes_equal(n1, n2):
348            if n1 is None and n2 is None:
349                return True
350            if n1 is None or n2 is None:
351                return False
352            return (
353                n1.value == n2.value and
354                nodes_equal(n1.left, n2.left) and
355                nodes_equal(n1.right, n2.right)
356            )
357        return nodes_equal(self.root, other.root)