refactor

Author: Brunods1001

ElitePython/_modidx.py

@@ -24,6 +24,12 @@
                                                                                     'ElitePython/binary_trees.py'),
                                           'ElitePython.binary_trees.bfs': ( 'Data Structures/Binary Trees/binary_trees.html#bfs',
                                                                             'ElitePython/binary_trees.py'),
+                                          'ElitePython.binary_trees.bst_insert': ( 'Data Structures/Binary '
+                                                                                   'Trees/binary_trees.html#bst_insert',
+                                                                                   'ElitePython/binary_trees.py'),
+                                          'ElitePython.binary_trees.bst_search': ( 'Data Structures/Binary '
+                                                                                   'Trees/binary_trees.html#bst_search',
+                                                                                   'ElitePython/binary_trees.py'),
                                           'ElitePython.binary_trees.create_n_nodes': ( 'Data Structures/Binary '
                                                                                        'Trees/binary_trees.html#create_n_nodes',
                                                                                        'ElitePython/binary_trees.py'),

ElitePython/binary_trees.py

@@ -1,86 +1,96 @@
-"""Create and operate on Binary Trees"""
-
 # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/Data Structures/Binary Trees/binary_trees.ipynb.
 
 # %% auto 0
 __all__ = ['visualize_tree', 'Node', 'set_parents', 'create_n_nodes', 'dfs_inorder', 'dfs_preorder', 'dfs_postorder', 'bfs',
-           'insert_node', 'search_tree', 'delete_node']
+           'insert_node', 'search_tree', 'delete_node', 'bst_insert', 'bst_search']
 
 # %% ../nbs/Data Structures/Binary Trees/binary_trees.ipynb 3
 from collections import deque
 import matplotlib.pyplot as plt
 import networkx as nx
 from typing import Callable, Generator
 
-# %% ../nbs/Data Structures/Binary Trees/binary_trees.ipynb 5
-def visualize_tree(root):
+# %% ../nbs/Data Structures/Binary Trees/binary_trees.ipynb 12
+def visualize_tree(root, show_parents=False):
+    """Visualize a binary tree using NetworkX and Matplotlib.
+    Args:
+        root: The root node of the tree.
+        show_parents: If True, include parent values in node labels.
+    """
     if not root:
         print("Empty tree")
         return
-    
+    import math
     G = nx.DiGraph()
     pos = {}
     labels = {}
     node_id = 0
-    
+    node_map = {}
     def add_nodes(node, x=0.0, y=0.0, layer=1):
         nonlocal node_id
         if node:
             current_id = node_id
+            node_map[node] = current_id
             node_id += 1
             G.add_node(current_id)
             pos[current_id] = (x, y)
-            labels[current_id] = str(node.value)
-            
+            if show_parents and hasattr(node, 'parent') and node.parent:
+                labels[current_id] = f"{node.value}\n(p:{node.parent.value})"
+            else:
+                labels[current_id] = str(node.value)
             if node.left:
                 left_id = node_id
                 G.add_edge(current_id, left_id)
-                add_nodes(node.left, x - 1/layer, y - 1, layer + 1)
+                add_nodes(node.left, x - 1/(layer+0.5), y - 1, layer + 1)
             if node.right:
                 right_id = node_id
                 G.add_edge(current_id, right_id)
-                add_nodes(node.right, x + 1/layer, y - 1, layer + 1)
-            
+                add_nodes(node.right, x + 1/(layer+0.5), y - 1, layer + 1)
             return current_id
-    
     add_nodes(root)
-    
-    plt.figure(figsize=(3, 2))
-    nx.draw(G, pos, labels=labels, with_labels=True, node_color='lightblue', 
-            node_size=1000, font_size=16, font_weight='bold', arrows=True)
+    # Adaptive figure size
+    depth = root.height if hasattr(root, 'height') else 3
+    width = max(6, 2 ** (depth-1))
+    plt.figure(figsize=(width, depth*2))
+    # Highlight root
+    node_colors = ["#ffcc00" if n == 0 else "#90caf9" for n in G.nodes]
+    nx.draw(G, pos, labels=labels, with_labels=True, node_color=node_colors,
+            node_size=1200, font_size=14, font_weight='bold', arrows=True, edge_color="#888888")
     plt.title("Binary Tree Visualization")
+    plt.axis('off')
     plt.show()
 
 # %% ../nbs/Data Structures/Binary Trees/binary_trees.ipynb 14
 class Node:
-    """A node in a binary tree."""
-    def __init__(self, key, left: "Node|None"=None, right: "Node|None"=None, parent: "Node|None" = None):
-        self.value = key
+    """A node in a binary tree with parent pointer."""
+    __slots__ = ['value', 'parent', 'left', 'right']
+
+    def __init__(self, value, left: "Node|None"=None, right: "Node|None"=None, parent: "Node|None"=None):
+        self.value = value
         self.parent: Node|None = parent
         self.left: Node|None = left
         self.right: Node|None = right
 
     def __str__(self) -> str:
-        """A string representation of the node including its value, left, and right children."""
+        """A string representation of the node including its value, left, right, and parent."""
         left_value = f"Node({self.left.value})" if self.left else None
         right_value = f"Node({self.right.value})" if self.right else None
         parent_value = f"Node({self.parent.value})" if self.parent else None
         return f"Node({self.value}, left={left_value}, right={right_value}, parent={parent_value})"
 
     def __repr__(self) -> str:
-        """A representation of the node including its value, left, and right children."""
         return self.__str__()
 
     @property
     def height(self) -> int:
+        """Returns the height of the node (longest path to a leaf)."""
         return 1 + max(
             (self.left.height if self.left else 0), (self.right.height if self.right else 0)
         )
 
     @property
     def depth(self) -> int:
-        return 1 + (self.parent.depth if self.parent else 0)
-
+        return 0 if self.parent is None else 1 + self.parent.depth
 
 # %% ../nbs/Data Structures/Binary Trees/binary_trees.ipynb 15
 def set_parents(node: Node, parent: Node|None=None) -> Node:
@@ -95,6 +105,7 @@ def set_parents(node: Node, parent: Node|None=None) -> Node:
 
 # %% ../nbs/Data Structures/Binary Trees/binary_trees.ipynb 21
 def create_n_nodes(num=5) -> Node:
+    """Creates a perfect binary tree with num nodes."""
     assert num > 0
     nodes = [Node(i) for i in range(num)]
     for i in range(num):
@@ -108,16 +119,9 @@ def create_n_nodes(num=5) -> Node:
             nodes[right_index].parent = nodes[i]
     return nodes[0]  # root node
 
-visualize_tree(create_n_nodes(1))
-visualize_tree(create_n_nodes(5))
-visualize_tree(create_n_nodes(9))
-
 # %% ../nbs/Data Structures/Binary Trees/binary_trees.ipynb 24
 def dfs_inorder(node: Node) -> Generator[Node, None, None]:
-    """Depth-first search.
-    
-    Visits left subtree, then node, then right subtree
-    """
+    """Inorder DFS traversal: left, node, right."""
     if node.left:
         yield from dfs_inorder(node.left)
     yield node
@@ -128,10 +132,7 @@ def dfs_inorder(node: Node) -> Generator[Node, None, None]:
 
 # %% ../nbs/Data Structures/Binary Trees/binary_trees.ipynb 25
 def dfs_preorder(node: Node) -> Generator[Node, None, None]:
-    """Depth-first search.
-
-    Visits the node, then left subtree, then right subtree.
-    """
+    """Preorder DFS traversal: node, left, right."""
     yield node
     if node.left:
         yield from dfs_preorder(node.left)
@@ -141,10 +142,7 @@ def dfs_preorder(node: Node) -> Generator[Node, None, None]:
 
 # %% ../nbs/Data Structures/Binary Trees/binary_trees.ipynb 26
 def dfs_postorder(node: Node) -> Generator[Node, None, None]:
-    """Depth-first search.
-
-    Visits the left subtree, then right subtree, then the node
-    """
+    """Postorder DFS traversal: left, right, node."""
     if node.left:
         yield from dfs_postorder(node.left)
     if node.right:
@@ -154,7 +152,7 @@ def dfs_postorder(node: Node) -> Generator[Node, None, None]:
 
 # %% ../nbs/Data Structures/Binary Trees/binary_trees.ipynb 28
 def bfs(node: Node) -> Generator[Node, None, None]:
-    """Breadth-first search."""
+    """Breadth-first search (level order traversal)."""
     queue = deque([node])
     while queue:
         node = queue.popleft()
@@ -166,45 +164,40 @@ def bfs(node: Node) -> Generator[Node, None, None]:
 list(bfs(root))
 
 # %% ../nbs/Data Structures/Binary Trees/binary_trees.ipynb 30
-def insert_node(root: Node, key):
-    """"""
+def insert_node(root: Node, value):
+    """Level-order BFS insert. Inserts a new node with the given value."""
     queue = deque([root])
     while queue:
         node = queue.popleft()
         if node.left:
             queue.append(node.left)
         else:
-            node.left = Node(key=key, left=None, right=None, parent=node)
+            node.left = Node(value=value, left=None, right=None, parent=node)
             break
-
         if node.right:
             queue.append(node.right)
         else:
-            node.right = Node(key=key, left=None, right=None, parent=node)
+            node.right = Node(value=value, left=None, right=None, parent=node)
             break
-    
 
 # %% ../nbs/Data Structures/Binary Trees/binary_trees.ipynb 33
-def search_tree(root: Node, key, traversal_func: Callable[[Node], Generator[Node, None, None]]) -> Node|None:
+def search_tree(root: Node, value, traversal_func: Callable[[Node], Generator[Node, None, None]]) -> Node|None:
+    """Search for a node with the given value using the specified traversal function."""
     for node in traversal_func(root):
-        if node.value == key:
+        if node.value == value:
             return node
     else:
         return None
 
 # %% ../nbs/Data Structures/Binary Trees/binary_trees.ipynb 36
-def delete_node(root: Node, key=None) -> Node|None:
-    """Uses BFS to delete the node with a key and swap it with the right-most node."""
-
+def delete_node(root: Node, value=None) -> Node|None:
+    """Deletes the node with the given value by swapping with the right-most node."""
     # if root has no children
-    if root.left or root.right:
-        pass
-    else:
+    if not (root.left or root.right):
         del root
         return None
-
-    # find the deepest root node and its parent
-    queue: deque[tuple[Node, Node|None]] = deque([(root, None)])  # root has no parent
+    # find the deepest node and its parent
+    queue: deque[tuple[Node, Node|None]] = deque([(root, None)])
     last_parent = None
     last_node = None
     while queue:
@@ -213,7 +206,6 @@ def delete_node(root: Node, key=None) -> Node|None:
             queue.append((last_node.left, last_node))
         if last_node.right:
             queue.append((last_node.right, last_node))
-
     # delete the leaf node
     if last_parent:
         if last_parent.right:
@@ -222,26 +214,22 @@ def delete_node(root: Node, key=None) -> Node|None:
             last_parent.left = None
     else:
         raise NotImplementedError("Parent of right-most node not found.")
-    
-    # find the node with value == key
-    if last_node and root.value == key:
+    # find the node with value == value
+    if last_node and root.value == value:
         last_node.left = root.left
         last_node.right = root.right
         del root
         return last_node
-    elif last_node:
-        pass
-    else:
+    elif not last_node:
         raise NotImplementedError
-
     target_queue: deque[Node] = deque([root])
     while target_queue:
         node = target_queue.popleft()
         if node.left:
             if node.left is last_node:
                 node.left = None
                 return root
-            elif node.left.value == key:
+            elif node.left.value == value:
                 node.left = last_node
                 last_node.parent = node
                 return root
@@ -250,9 +238,41 @@ def delete_node(root: Node, key=None) -> Node|None:
             if node.right is last_node:
                 node.right = None
                 return root
-            elif node.right.value == key:
+            elif node.right.value == value:
                 node.right = last_node
                 last_node.parent = node
                 return root
             target_queue.append(node.right)
     return root
+
+# %% ../nbs/Data Structures/Binary Trees/binary_trees.ipynb 41
+def bst_insert(root: Node, new_node: Node) -> None:
+    """Insert a new node into a BST. Raises an exception if value exists."""
+    if new_node.value < root.value:
+        if root.left:
+            bst_insert(root.left, new_node)
+        else:
+            root.left = new_node
+            new_node.parent = root
+            return
+    elif root.value < new_node.value:
+        if root.right:
+            bst_insert(root.right, new_node)
+        else:
+            root.right = new_node
+            new_node.parent = root
+            return
+    else:
+        raise Exception(f"Value {root.value} already exists in this tree.")
+
+# %% ../nbs/Data Structures/Binary Trees/binary_trees.ipynb 45
+def bst_search(root: Node|None, key) -> Node:
+    """Search for a key in a BST. Raises an exception if not found."""
+    if root.value == key:
+        return root
+    elif key < root.value and root.left:
+        return bst_search(root.left, key)
+    elif root.right:
+        return bst_search(root.right, key)
+    else:
+        raise Exception(f"Key {key} not found")