STD_2 - Binary Search Tree.py

""" Basic BST code for inserting (i.e. building) and printing a tree

    Your ***second standard viva task*** (of 5) will be to implement a find method into
    the class BinaryTree from pseudocode. See the lab task sheet for Week 5. 

    Your ***first advanced viva task*** (of 3) will be to implement a remove (delete) method
    into the class Binary Tree from partial pseudocode. See the lab task sheet for Week 5 (available in Week 5).

    There will be some ***introductory challenges*** in Week 4, with solutions released in Week 5.
    It is STRONGLY RECOMMENDED you attempt these!

    Since the given code is in python it is strongly suggested you stay with python; but
    if you want to reimplement as C++ this is also OK (see the Week 5 lab sheet guidance). 
"""

import math

""" Node class
"""

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

""" BST class with insert and display methods. display pretty prints the tree
"""

class BinaryTree:
    def __init__(self):
        self.root = None
#v v v v v v v v v v v v v v v v v v v v v v v v v v v v v v v v v v v v v

    '''Binary Search Tree
       input: Values in a tree
       output: Target Value
       Implements a search function for a binary tree
    '''

    def find_i(self, data):             #pass in values defined elsewhere
        print("iteratively searching for: ", data)  #tell the user what the function is doing
        cur_node = self.root            #set the current node to be equal to the root value
        while cur_node != None:         #loop while there is a node to check
            if cur_node.data == data:   #if the current node is equal to the value being searched for
                print("Value found")
                return True                 #target value is found, return True
            elif cur_node.data > data:      #if the current node is greater than the target value
                cur_node = cur_node.left    #set the current node as the left node
                print("Current node is greater than target, moving to the left child")
            else:                           #if the current node is smaller than the target
                cur_node = cur_node.right   #set the current node as the right node
                print("Current node is less than the target, moving to the right child")
        return False                        #target value isnt found, return false, this whole section is looped until the value is found
    
#####################
    
    def find_r(self, data):             #pass in values defined elsewhere
        print("recursively searching for: ", data)  #tell the user what the function is doing
        return self._find_r(data, self.root)        #pass data to the main search function

    def _find_r(self, data, cur_node):
        if cur_node.data == data:       #if the current node is equal to the value being searched for
            print("Value found")
            return True                 #target value is found, return True
        elif cur_node.data > data:      #if the current node is greater than the target value
            print("Current node is greater than target, moving to the left child")
            return self._find_r(data,cur_node.left)     #recursively call the function
        else:                           #if the current node is smaller than target
            print("Current node is less than the target, moving to the right child")
            return self._find_r(data,cur_node.right)    #recursively call the function
        return False                    #target value isnt found
    
#^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^
    def insert(self, data):
        if self.root is None:
            self.root = Node(data)
        else:
            self._insert(data, self.root)

    def _insert(self, data, cur_node):
        if data < cur_node.data:
            if cur_node.left is None:
                cur_node.left = Node(data)
            else:
                self._insert(data, cur_node.left)
        elif data > cur_node.data:
            if cur_node.right is None:
                cur_node.right = Node(data)
            else:
                self._insert(data, cur_node.right)
        else:
            print("Value already present in tree")

    def display(self, cur_node):
        lines, _, _, _ = self._display(cur_node)
        for line in lines:
            print(line)


    def _display(self, cur_node):
        
        if cur_node.right is None and cur_node.left is None:
            line = '%s' % cur_node.data
            width = len(line)
            height = 1
            middle = width // 2
            return [line], width, height, middle

        if cur_node.right is None:
            lines, n, p, x = self._display(cur_node.left)
            s = '%s' % cur_node.data
            u = len(s)
            first_line = (x + 1) * ' ' + (n - x - 1) * '_' + s
            second_line = x * ' ' + '/' + (n - x - 1 + u) * ' '
            shifted_lines = [line + u * ' ' for line in lines]
            return [first_line, second_line] + shifted_lines, n + u, p + 2, n + u // 2
        
        if cur_node.left is None:
            lines, n, p, x = self._display(cur_node.right)
            s = '%s' % cur_node.data
            u = len(s)
            first_line = s + x * '_' + (n - x) * ' '
            second_line = (u + x) * ' ' + '\\' + (n - x - 1) * ' '
            shifted_lines = [u * ' ' + line for line in lines]
            return [first_line, second_line] + shifted_lines, n + u, p + 2, u // 2

        left, n, p, x = self._display(cur_node.left)
        right, m, q, y = self._display(cur_node.right)
        s = '%s' % cur_node.data
        u = len(s)
        first_line = (x + 1) * ' ' + (n - x - 1) * '_' + s + y * '_' + (m - y) * ' '
        second_line = x * ' ' + '/' + (n - x - 1 + u + y) * ' ' + '\\' + (m - y - 1) * ' '
        if p < q:
            left += [n * ' '] * (q - p)
        elif q < p:
            right += [m * ' '] * (p - q)
        zipped_lines = zip(left, right)
        lines = [first_line, second_line] + [a + u * ' ' + b for a, b in zipped_lines]
        return lines, n + m + u, max(p, q) + 2, n + u // 2

#example calls, which construct and display the tree       
bst = BinaryTree()
bst.insert(4)
bst.insert(2)
bst.insert(6)
bst.insert(1)
bst.insert(3)
bst.insert(5)
bst.insert(7)
bst.insert(8)
#bst.insert(9)
#bst.insert(10)
#bst.insert(11)
#bst.insert(12)
#bst.insert(13)
#bst.insert(14)
#bst.insert(15)
#bst.insert(100)
#bst.insert(200)


bst.find_i(2)
print()
bst.find_r(8)

bst.display(bst.root)