BST.py

import BSTExceptions

class Node(object):
    ''' A class that will hold a node in our Binary search tree '''
    def __init__(self, value):
        self.value=value # String value of the word
        self.frequency=1 # How many times the word is inserted in the tree
        self.left=None # Left child
        self.right=None # Right child
        self.parent=None # Parent node

    def insert(self, value):
        return tree_insert(self, value)

    def delete(self, value):
        # We need to replace self with the new object returned from tree_delete function
        # just doing self = newObj won't work, so we use a workaround updating
        # self.__dict__, which is a dictionary that contains all of the object's attributes
        newObj = tree_delete(self, value)
        self.__dict__.update(newObj.__dict__)
        return newObj

    def pre_order(self):
        return pre_order(self)

    def in_order(self):
        return in_order(self)
    
    def count_children(self):
        return count_children(self)

    def print(self):
        tree_print(self)
        return self

def tree_insert(tree, item): # O(n)
    ''' Inserts a new node into the tree or updates the frequency, if the node already exists '''
    if tree==None:
        # If our tree is empty, initialize it with the provided item as root
        tree=Node(item)
    else:
        if(item == tree.value):
            # If item already in our binary search tree, increase it's frequency
            tree.frequency += 1
        elif(item < tree.value):
            # If item alphabetically comes before, go left
            if(tree.left==None):
                # If there is no left child, we insert our item as the left child
                tree.left=Node(item)
                tree.left.parent = tree
            else:
                # If left child exists, run insert function again with left child as root
                tree_insert(tree.left, item)
        else:
            # If item alphabetically comes after, go right
            if(tree.right==None):
                # If there is no right child, we insert our item as the right child
                tree.right=Node(item)
                tree.right.parent = tree
            else:
                # If right child exists, run insert function again with right child as root
                tree_insert(tree.right, item)
    return tree

def pre_order(tree): # O(n)
    ''' Prints the tree in pre order: Prints the element, then goes left, then right '''
    if(tree==None):
        raise BSTExceptions.NoTree
    out = [] # Apart from printing, generate a list of order
    tree_print(tree)
    out.append([tree.value,tree.frequency])
    if(tree.left!=None):
        out += pre_order(tree.left)
    if(tree.right!=None):
        out += pre_order(tree.right)
    return out

def in_order(tree): # O(n)
    ''' Prints the tree in alphabetical order '''
    if(tree==None):
        raise BSTExceptions.NoTree
    out = [] # Apart from printing, generate a list of order
    if(tree.left != None):
        out += in_order(tree.left)
    tree_print(tree)
    out.append(tree.value)
    if(tree.right != None):
        out += in_order(tree.right)
    return out

def tree_print(tree): # O(1)
    ''' An additional function, that prints elements with information 
     about their parent and frequency '''
    if(tree==None):
        raise BSTExceptions.NoTree
    parent = tree.parent
    if(tree.parent==None):
        parent = "root"
    else:
        parent = "parent: " + tree.parent.value
    print(tree.value, tree.frequency, parent)

def tree_find(tree, target): # O(n)
    ''' Find if the target value exists in the tree and displays the route
     used to find the target. When found, prints out frequency of the target '''
    if(tree == None):
        # Tree is empty, so target is not in the tree
        print(target,"not found")
        return tree
    elif(tree.value==target):
        # Target has been found
        print(target,"found with frequency:",tree.frequency)
        return tree
    elif(target<tree.value):
        # Target alphabetically lower than root, call function again
        # with left child as root
        print("Looking at:", tree.value,", going left")
        return tree_find(tree.left, target)
    else:
        # Target alphabetically higher than root, call function again
        # with right child as root
        print("Looking at:", tree.value, ", going right")
        return tree_find(tree.right, target)
    

def count_children(node): # O(1)
    ''' Counts if a node has 0, 1 or 2 children '''
    if(node==None):
        raise BSTExceptions.NoTree
    count = 0
    if(node.left!=None):
        count += 1
    if(node.right!=None):
        count += 1
    return count

def tree_delete(tree, target): # O(n)
    ''' Find target node in the tree '''
    # Raise exception if tree is None
    if(tree==None):
        raise BSTExceptions.NoTree

    target_node = tree_find(tree, target)
    if(target_node==None):
        # Target node is not in the tree
        print(target,"is not in this tree")
        return tree
    # Count the number of children for target node
    children = count_children(target_node)
    
    if(children==0):
        # Since no children, we remove the node and remove parent connection
        if(target_node.parent!=None):
            if(node_to_left_of_parent(target_node)):
                target_node.parent.left = None
            else:
                target_node.parent.right = None
        else:
            # Deleted node was root
            tree = None
        del target_node
        print(target,"deleted from the tree")
        return tree

    if(children==1):
        # If only one child, we interchange the child node with the parent node
        # and remove the current node
        child = target_node.left
        if(child == None):
            child = target_node.right
        child.parent = target_node.parent
        # Figure out if our target is to the left or to the right of its parent
        if(target_node.parent!=None):
            if(node_to_left_of_parent(target_node)):
                # It is to the left
                target_node.parent.left = child
            else:
                # It is to the right
                target_node.parent.right = child
        else:
            # Deleted node was root, so set root as child
            tree = child
        del target_node
        print(target,"deleted from the tree")
        return tree

    else:
        # Our node has two children. Find the max value from left subtree
        # and swap our target node with the new found node. Remove the
        # duplicate node
        max_node = max_value_node(target_node.left)
        target_node.value = max_node.value
        target_node.frequency = max_node.frequency
        print(target,"swapped with",max_node.value)
        print("Now removing duplicate:",max_node.value)
        tree_delete(target_node.left, max_node.value)
        return tree


def max_value_node(node): # O(n)
    ''' Finds node with the highest alphabetical value in provided subtree '''
    # Maximum value will be on the far right hand child, so we go right
    # until we don't have a right child
    if(node==None):
        raise BSTExceptions.NoTree
    if(node.right==None):
        return node
    return max_value_node(node.right)

def node_to_left_of_parent(node): # O(1)
    ''' Helper function that determines if provided node is the left child
     for its parent node. Returns True if is the left child and False if
     is the right child '''
    # We include less than or equal to, because in scenario where we have a tree:
    #     B
    #   A   C
    # and A has been swapped with B because of our delete function leading to:
    #     B
    #   B   C
    # the second B is still to left of first B
    if(node==None):
        raise BSTExceptions.NoTree
    if(node.value <= node.parent.value):
        return True
    else:
        return False

if __name__ == "__main__":
    tree = None
    # We remove punctuation in the most efficient way using str.translate()
    punctuation = '!"#$%&\'()*+,-./:;<=>?@[\\]^_`{|}~'
    table = str.maketrans({key: None for key in punctuation})
    with open('words.txt','r') as f:
        for line in f:
            line = line.translate(table)
            for word in line.split():
                tree = tree_insert(tree, word)
    tree.pre_order()
    print()
    tree.delete("yet")
    print()
    tree.pre_order()
    print()
    tree_find(tree, "to")