[Data Structure] Python: Binary Search Tree

A binary search tree can be implemented in Python using a Node class and a BinarySearchTree class. 

The Node class contains a key value and data members for the left and right children nodes. 

The BinarySearchTree class contains a data member for the tree's root (a Node object).

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

class BinarySearchTree:
    def __init__(self):
    	self.root = None

 

 

Search() Method

The BST search algorithm first assigns current_node with the root. 

def search(self, desired_key):
    current_node = self.root
    while current_node is not None:
    	# Return the node if the key matches. 
        if current_node.key == desired_key:
        	return current_node
        
        # Navigate to the left if the search key is 
        # less than the node's key
        elif desired_key < current_node.key:
        	current_node = currnet_node.left
        
        # Navigate to the right if the search key is
        # greater than the node's key
        else:
        	current_node = current_node.right

    # Key not found 
    return None

 

 

Insert() Method

The insert() method is used to insert a new node into the tree.

If the tree is empty, then the root data member is assigned with the new node.

If the tree is not empty, current_node is assigned with the root node, and a loop is entered. 

 

Inside the loop, if the new node's key is less than the current node's key, then the new node is inserted as the current node's left child (if the current node has no chidl), or current_node is assigned with current_node.left. 

If the new node's key is greater than or equal to the current node's key, then the new node is inserted as the current node's right child (if the current node has no right child), or current_node is assigned with current_node.right. 

def insert(self, node):
    # Check if the tree is empty
    if self.root is None:
    	self.root = node
    else:
    	current_node = self.root
        while current_node is not None:
        	if node.key < current_node.key:
            	# If there is no left child, add the new node here
                # Otherwise repeat from the left child.
                if current_node.left is None:
                	current_node.left = node
                    current_node = None
                else:
	                current_node = current_node.left
            else:
            	# If there is no right child, add the new node here
                # Otherwise repeat from the right child.
                if current_node.right is None:
                	current_node.right = node
                    current_node = None
                else:
                	current_node = current_node.right

 

 

Remove() Method

The remove() method is more complicated than serach() or insert() because several cases exist. 

Different actions are taken depending on what kind of node is being removed. 

• Case 1: The node being removed is a leaf. The parent's left or right data member is assigned with None, depending on which side the node is.
• Case 2: The node being removed has one child. The parent's left or right data member is assigned with the removed node's single child.
• Case 3:The node being removed has two children. The node's key is replaced by the successor's key, and then the successor (which will fall under either Case 1 or Case 2) is removed. The successor is the next largest node in the tree, and is always the right child's leftmost child.

def remove(self, key):
    parent = None
    current_node = self.root
    
    # Search for the node.
    while current_node is not None:
    
        # Check if current_node has a matching key.
        if current_node.key == key: 
            if current_node.left is None and current_node.right is None:   # Case 1
                if parent is None: # Node is root
                    self.root = None
                elif parent.left is current_node: 
                    parent.left = None
                else:
                    parent.right = None
                return  # Node found and removed
            elif current_node.left is not None and current_node.right is None:  # Case 2
                if parent is None: # Node is root
                    self.root = current_node.left
                elif parent.left is current_node: 
                    parent.left = current_node.left
                else:
                    parent.right = current_node.left
                return  # Node found and removed
            elif current_node.left is None and current_node.right is not None:  # Case 2
                if parent is None: # Node is root
                    self.root = current_node.right
                elif parent.left is current_node:
                    parent.left = current_node.right
                else:
                    parent.right = current_node.right
                return  # Node found and removed
            else:                                    # Case 3
                # Find successor (leftmost child of right subtree)
                successor = current_node.right
                while successor.left is not None:
                    successor = successor.left
                current_node.key = successor.key      # Copy successor to current node
                parent = current_node
                current_node = current_node.right     # Remove successor from right subtree
                key = parent.key                      # Loop continues with new key
        elif current_node.key < key: # Search right
            parent = current_node
            current_node = current_node.right
        else:                        # Search left
            parent = current_node
            current_node = current_node.left
            
    return # Node not found