BinarySearchTree.py

from BinaryTree import BinaryTree
from Interfaces import Set


class BinarySearchTree(BinaryTree, Set):

    def __init__(self):
        BinaryTree.__init__(self)
        self.n = 0

    def add(self, key: object, value: object = None) -> bool:
        """
        If the key does not exist in this BinarySearchTree,
        adds a new node with given key and value, in the correct position.
        Returns True if the key-value pair was added to the tree, False otherwise.
        """
        new = self.Node(key, value)
        parent = self._find_last(key)
        child = self._add_child(parent, new)
        return child

    def find(self, key: object) -> object:
        """
        returns the value corresponding to the given key if the key
        exists in the BinarySearchTree, None otherwise
        """
        node = self._find_eq(key)
        if node is None:
            return None
        return node.v

    def remove(self, key: object):
        """
        removes the node with given key if it exists in this BinarySearchTree.
        Returns the value corresponding to the removed key, if the key was in the tree.
        If given key does not exist in the tree, ValueError is raised.
        """
        u = self._find_eq(key)
        if u is None:
            raise ValueError("Key is not in tree")
        value = u.v
        self._remove_node(u)
        return value

    def _find_eq(self, key: object) -> BinaryTree.Node:
        """
        helper method; returns the node in this tree that contains the given key,
        None otherwise.
        """
        current = self.r
        while current is not None:
            if key < current.k:
                current = current.left
            elif key > current.k:
                current = current.right
            else:
                return current  # returns key when it is found
        return current  # if no key matches, returns None

    def _find_last(self, key: object) -> BinaryTree.Node:
        """
        helper method; returns the node in this tree that contains the given key, if it exists.
        Otherwise, returns the node that would have been the parent of the node
        with the given key, if it existed
        """
        current = self.r
        parent = None
        while current is not None:
            parent = current
            if key < current.k:
                current = current.left
            elif key > current.k:
                current = current.right
            else:
                return current  # returns key when it is found
        return parent  # returns parent if key is not found

    def _add_child(self, p: BinaryTree.Node, u: BinaryTree.Node) -> bool:
        """
        helper method; adds node u as the child of node p, assuming node p has at most 1 child
        """
        if p is None:  # if p is nil, make u the root
            self.r = u  # insert into empty tree
        else:  # determines position of u
            if u.k < p.k:
                p.left = u
            elif u.k > p.k:
                p.right = u
            else:
                return False  # if key already exists
            u.parent = p
        self.n += 1
        return True

    def _splice(self, u: BinaryTree.Node):
        """
        helper method; links the parent of given node u to the child
        of node u, assuming u has at most one child
        """
        if u.left is not None:
            child = u.left
        else:
            child = u.right
        if u == self.r:
            self.r = child
            p = None
        else:
            p = u.parent
            if p.left == u:
                p.left = child
            else:
                p.right = child
        if child is not None:
            child.parent = p
        self.n -= 1

    def _remove_node(self, u: BinaryTree.Node):
        if u.left is None or u.right is None:  # checks if u has 0 or 1 child, splices u
            self._splice(u)
        else:  # finds key with the largest value, then replaces keys and values
            w = u.right
            while w.left is not None:
                w = w.left
            u.k = w.k
            u.v = w.v
            self._splice(w)

    def find_smallest_right_node(self, key: object):  # finds the smallest node on the right side
        current = self.r
        smallest = None
        while current is not None:
            if key < current.k:
                smallest = current
                current = current.left
            elif key > current.k:
                current = current.right
            else:
                return current
        return smallest

    def clear(self):
        """
        empties this BinarySearchTree
        """
        self.r = None
        self.n = 0

    def __iter__(self):
        u = self.first_node()
        while u is not None:
            yield u.k
            u = self.next_node(u)

    def first_node(self):
        w = self.r
        if w is None: return None
        while w.left is not None:
            w = w.left
        return w

    def next_node(self, w):
        if w.right is not None:
            w = w.right
            while w.left is not None:
                w = w.left
        else:
            while w.parent is not None and w.parent.left != w:
                w = w.parent
            w = w.parent
        return w