binary search tree

binary_search_tree.py

#coding:UTF-8

'''
notes on R5

What is a data structure?
Algorithms for you to query update store information

BST
query:
    * min - go left
    * max - go right
    * search - < node THEN go left or > node THEN go right
    * next larger
    * next smaller
update:
    * insert - similar to search
    * delete - 

* all operations - O(N) for worst case - O(lg(N)) for average case

What is the  RI(repesentative invariant)?

How to balance the tree to avoid the O(N) worst case?
'''

from functools import total_ordering

class BST:

    class SafeDeleteException(Exception): pass

    def __init__(self, root=None):
        self.root = root

    def min(self):
        node = self.root
        if node.left is None:
            return node
        else:
            return node.left.min()

    def max(self):
        node = self.root
        if node.right is None:
            return node
        else:
            return node.right.max()

    def insert(self, node):
        if self.root is None:
            self.root = node
            return 

        if node == self.root:
            raise Exception("{} already existed".format(node))

        elif node > self.root:
            if self.root.right is None:
                self.root.right = node
            else:
                self.root.right.insert(node)
        else:
            if self.root.left is None:
                self.root.left = node
            else:
                self.root.left.insert(node)

    def search(self, node):
        if node == self.root:
            return True, self.root

        if node > self.root:
            if self.root.right is None:
                return False, None
            else:
                return self.root.right.search(node)
        else:
            if self.root.left is None:
                return False, None
            else:
                return self.root.left.search(node)

    def next_larger(self, node):
        # node must be in the tree
        found, node_in_tree = self.search(node)
        if not found:
            raise Exception("{} not found in tree".format(node))

        if node_in_tree.right is not None:
            return node_in_tree.right.min()

        # find the 1st parent who is itself a left child
        # start from the node itself
        n = node_in_tree
        while not n.is_root() and not n.is_left_child():
            if n.is_root():
                raise Exception("{} is min, no next_larger".format(node_in_tree))
            n = n._parent
        return n._parent

    def next_smaller(self, node):
        found, node_in_tree = self.search(node)
        if not found:
            raise Exception("{} not found in tree".format(node))

        if node_in_tree.left is not None:
            return node_in_tree.left.max()

        n = node_in_tree
        while not n.is_right_child():
            if n.is_root():
                raise Exception("{} is min, no next_larger".format(node_in_tree))
            n = n._parent
        return n._parent

    def delete(self, node):
        found, node_in_tree = self.search(node)
        if not found:
            raise Exception("{} not found in tree".format(node))

        try:
            self.safe_delete(node_in_tree)
        except BST.SafeDeleteException:
            # node_in_tree has both left child and right child
            next_larger_node = self.next_larger(node_in_tree)
            next_larger_node._val, node_in_tree._val = node_in_tree._val, next_larger_node._val
            node_in_tree.right.safe_delete(next_larger_node)


    def safe_delete(self, node_in_tree):
        # this part is really ugly
        # TODO: simplify it!!!
        if node_in_tree.is_leaf():
            if node_in_tree.is_left_child():
                node_in_tree._parent.left = None
            elif node_in_tree.is_right_child():
                node_in_tree._parent.right = None
            elif node_in_tree.is_root():
                self = BST()
            else:
                assert False, "not possible"
        elif node_in_tree.has_left_child() and not node_in_tree.has_right_child():
            if node_in_tree.is_left_child():
                node_in_tree._parent._left = node_in_tree._left
                node_in_tree._left.root._parent = node_in_tree._parent
                node_in_tree._left.root._side = node_in_tree._side
            elif node_in_tree.is_right_child():
                node_in_tree._parent._right = node_in_tree._left
                node_in_tree._left.root._parent = node_in_tree._parent
                node_in_tree._left.root._side = node_in_tree._side
            else:
                assert False, "not possible"
        elif node_in_tree.has_right_child() and not node_in_tree.has_left_child():
            if node_in_tree.is_left_child():
                node_in_tree._parent._left = node_in_tree._right
                node_in_tree._right.root._parent = node_in_tree._parent
                node_in_tree._right.root._side = node_in_tree._side
            elif node_in_tree.is_right_child():
                node_in_tree._parent._right = node_in_tree._right
                node_in_tree._right.root._parent = node_in_tree._parent
                node_in_tree._right.root._side = node_in_tree._side
            else:
                assert False, "not possible"
        else:
            raise BST.SafeDeleteException("this node has both left child and right child.")

@total_ordering
class Node:
    ROOT = 0
    LEFT = 1
    RIGHT = 2

    def __init__(self, val):
        self._val = val
        self._left = None
        self._right = None
        self._parent = None
        self._side = Node.ROOT

    def __eq__(self, node):
        return self._val == node._val

    def __gt__(self, node):
        return self._val > node._val

    @property
    def left(self):
        return self._left

    @left.setter
    def left(self, node):
        if node is None:
            self._left = None
        else:
            node._parent = self
            node._side = Node.LEFT
            self._left = BST(node)

    @property
    def right(self):
        return self._right

    @right.setter
    def right(self, node):
        if node is None:
            self._right = None
        else:
            node._parent = self
            node._side = Node.RIGHT
            self._right = BST(node)

    def is_root(self):
        return self._side == Node.ROOT

    def is_left_child(self):
        return self._side == Node.LEFT 

    def is_right_child(self):
        return self._side == Node.RIGHT

    def has_left_child(self):
        return self.left is not None

    def has_right_child(self):
        return self.right is not None

    def is_leaf(self):
        return self.left is None and self.right is None

    def __repr__(self):
        return "Node({})".format(self._val)



l = [5, 3, 8, 7, 9, 1, 2, 10, 4, 6]
bst = BST()
for i in l:
    bst.insert(Node(i))

# print(bst.max())
# print(bst.min())
# print(bst.next_larger(Node(1)))
# print(bst.next_larger(Node(2)))
# print(bst.next_larger(Node(3)))
# print(bst.next_larger(Node(5)))
# print(bst.next_larger(Node(7)))
# print(bst.next_larger(Node(8)))
# # print(bst.next_larger(Node(9)))



# # print(bst.next_smaller(Node(1)))
# print(bst.next_smaller(Node(2)))
# print(bst.next_smaller(Node(3)))
# print(bst.next_smaller(Node(5)))
# print(bst.next_smaller(Node(7)))
# print(bst.next_smaller(Node(8)))
# print(bst.next_smaller(Node(9)))

# assert bst.next_larger(Node(2)) == Node(3)
# bst.delete(Node(3))
# assert bst.next_larger(Node(2)) == Node(4)
# bst.delete(Node(4))
# assert bst.next_larger(Node(2)) == Node(5)

assert bst.next_larger(Node(7)) == Node(8)
bst.delete(Node(8))
bst.delete(Node(9))
assert bst.next_larger(Node(7)) == Node(10)
bst.delete(Node(10))
assert bst.next_larger(Node(5)) == Node(6)

bst.delete(Node(4))
assert bst.next_smaller(Node(5)) == Node(3)