Python3 simple BST

bst.py

#! env python3

class Node:
    def __init__(self, value):
        self.value = value
        self.left = None
        self.right = None

    def insert(self, root, key):
        # If the root is None, create a new node with the key and return it
        if root is None:
            return Node(key)
        else:
            # If the key is already present in the tree, return the root
            if root.value == key:
                return root
            # If the key is greater than the root's value, insert it in the right subtree
            if root.value < key:
                root.right = self.insert(root.right, key)
            # If the key is less than the root's value, insert it in the left subtree
            else:
                root.left = self.insert(root.left, key)
        # Return the root after insertion
        return root
    
    def search(self, root, key):
        # Base case: return the node if it's found or None if it's not
        if root is None or root.value == key:
            return root
        # If the key is greater than the root's value, search in the right subtree
        if root.value < key:
            return self.search(root.right, key)
        # If the key is less than the root's value, search in the left subtree
        return self.search(root.left, key)
    
    def delete(self, root, key):
        # Base case: if the root is None, return None
        if root is None:
            return root
        # If the key to be deleted is smaller than the root's key,
        # then it lies in the left subtree
        if key < root.value:
            root.left = self.delete(root.left, key)
        # If the key to be deleted is greater than the root's key,
        # then it lies in the right subtree
        elif key > root.value:
            root.right = self.delete(root.right, key)
        # If the key is the same as the root's key, then this is the node
        # to be deleted
        else:
            # Node with only one child or no child
            if root.left is None:
                return root.right
            elif root.right is None:
                return root.left
            # Node with two children: Get the inorder successor
            # (smallest in the right subtree)
            temp = self.minValueNode(root.right)
            # Copy the inorder successor's content to this node
            root.value = temp.value
            # Delete the inorder successor
            root.right = self.delete(root.right, temp.value)
        return root

    def minValueNode(self, node):
        # Loop to find the leftmost leaf
        current = node
        while current.left is not None:
            current = current.left
        return current
    
    def balanceBST(self, root):
        # Helper function to store the in-order traversal of the BST
        def storeInOrder(node, inorder_nodes):
            if not node:
                return
            storeInOrder(node.left, inorder_nodes)
            inorder_nodes.append(node.value)
            storeInOrder(node.right, inorder_nodes)
        
        # Helper function to build a balanced BST from sorted nodes
        def buildBalancedBST(nodes, start, end):
            if start > end:
                return None
            mid = (start + end) // 2
            node = Node(nodes[mid])
            node.left = buildBalancedBST(nodes, start, mid - 1)
            node.right = buildBalancedBST(nodes, mid + 1, end)
            return node
        
        # Store the in-order traversal of the BST
        inorder_nodes = []
        storeInOrder(root, inorder_nodes)
        
        # Build and return the balanced BST
        return buildBalancedBST(inorder_nodes, 0, len(inorder_nodes) - 1)
    
    def __str__(self):
        lines = []
        self._print_tree(self, lines, "", "")
        return "\n".join(lines)

    def _print_tree(self, node, lines, prefix, children_prefix):
        if node is not None:
            lines.append(prefix + str(node.value))
            if node.left is not None or node.right is not None:
                if node.right is not None:
                    self._print_tree(node.right, lines, children_prefix + "├── ", children_prefix + "│   ")
                else:
                    lines.append(children_prefix + "├── ")
                if node.left is not None:
                    self._print_tree(node.left, lines, children_prefix + "└── ", children_prefix + "    ")
                else:
                    lines.append(children_prefix + "└── ")


tree = Node(10)
tree.insert(tree, 18)
tree.insert(tree, 15)
tree.insert(tree, 12)
tree.insert(tree, 8)
tree.insert(tree, 7)
tree.insert(tree, 6)
tree.insert(tree, 5)
tree.insert(tree, 4)
tree.insert(tree, 3)
tree.insert(tree, 2)

print(tree)

print(tree.balanceBST(tree))