BinarySearchTree.py

import random class NodeBST(object): def __init__(self, key=None, value=None): self.key = key self.value = value self.parent = None self.left_child = None self.right_child = None def __str__(self): return "".format(self.key, self.value) __repr__ = __str__ def reset(self): self.key = None self.value = None self.parent = None self.left_child = None self.right_child = None class BinarySearchTree(object): def __init__(self): self.root = None def __str__(self): return "" __repr__ = __str__ def __compare(self, key=None, method='compare', source=None, print_path=False): compare_node = source if source else self.root parent_node = None while compare_node: parent_node = compare_node if method == 'search': if parent_node.key == key: # when the method is search, compare_node is the result parent_node = parent_node.parent break if method == 'compare' or method == 'search': compare_node = parent_node.left_child if key <= parent_node.key else parent_node.right_child if method == 'min': compare_node = parent_node.left_child if method == 'max': compare_node = parent_node.right_child if print_path: try: print("({}, {}) -> ({}, {})".format(parent_node.key, parent_node.value, compare_node.key, compare_node.value)) except AttributeError: pass return parent_node, compare_node def __check_root(self): """ Check whether the tree is empty. """ if not self.root or not self.root.key: raise IndexError("The tree is empty!") def __check_node(self, node): """ Check whether a node exists. Args: node: class NodeBST """ if not node or not node.key: raise IndexError("Node doesn't exist!") def __swap_kv(self, node1, node2): """ Swap key-value pair of two node. Args: node1: class NodeBST node2: class NodeBST """ node1.key, node2.key = node2.key, node1.key node1.value, node2.value = node2.value, node1.value def rotation(self, key, right_rotation=False): """ Rotate around a certain node. Args: key: the key of node to be rotated right_rotation: True for right rotation, False for left rotation """ parent_node, node = self.__compare(key, method='search') # left rotation if not right_rotation: self.__check_node(node.right_child) # update parent parent = node.parent neighbor = node.right_child if node.key <= parent.key: parent.left_child = neighbor else: parent.right_child = node.right_child # update node node.parent = neighbor if neighbor.left_child and neighbor.left_child.key: node.right_child = neighbor.left_child # update neighbor neighbor.parent = parent neighbor.left_child = node # right rotation else: self.__check_node(node.left_child) # update parent parent = node.parent neighbor = node.left_child if node.key <= parent.key: parent.left_child = neighbor else: parent.right_child = node.right_child # update node node.parent = neighbor if neighbor.right_child and neighbor.right_child.key: node.left_child = neighbor.right_child # update neighbor neighbor.parent = parent neighbor.right_child = node def insert(self, key, value): """ Insert a key-value pair. """ insert_node = NodeBST(key, value) # insert root if not self.root or not self.root.key: self.root = insert_node else: # find the position to insert the node parent_node, _ = self.__compare(key) insert_node.parent = parent_node if key <= parent_node.key: parent_node.left_child = insert_node else: parent_node.right_child = insert_node print("Insert {} : {}".format(key, value)) def get_node(self, key, print_path=False): """ Get the node with the given key. Args: key: the key of the node print_path: True for printing the searching path, and vice versa Returns: search_node: class NodeBST """ self.__check_root() parent_node, search_node = self.__compare(key, method='search', print_path=print_path) self.__check_node(search_node) return search_node def search(self, key, print_path=False): """ Search for a certain key. Print the information of the node. Args: key: the key of the node to be searched print_path: True for printing the searching path, and vice versa """ search_node = self.get_node(key, print_path) print("ID: {}\nValue: {}".format(search_node.key, search_node.value)) @property def min(self): self.__check_root() min_node, _ = self.__compare(method='min') return min_node if min_node else None @property def max(self): self.__check_root() max_node, _ = self.__compare(method='max') return max_node if max_node else None def get_predecessor(self, key): """ Get the predecessor the of given node. Args: key: the key of the node to be searched Returns: pred_node: predecessor of the node, class NodeBST """ self.__check_root() parent_node, search_node = self.__compare(key, method='search') self.__check_node(search_node) # if the node has a left tree if search_node.left_child and search_node.left_child: pred_node, _ = self.__compare(method='max', source=search_node.left_child) # if the node has no left tree else: while search_node.key < parent_node.key: search_node = parent_node parent_node = parent_node.parent # if it reaches the root, means there is no predecessor if not parent_node: return NodeBST(None, None) pred_node = parent_node return pred_node def get_successor(self, key): """ Get the successor the of given node. Args: key: the key of the node to be searched Returns: succ_node: successor of the node, class NodeBST """ self.__check_root() parent_node, search_node = self.__compare(key, method='search') self.__check_node(search_node) if search_node.right_child and search_node.right_child.key: succ_node, _ = self.__compare(method='min', source=search_node.right_child) else: while search_node.key > parent_node.key: search_node = parent_node parent_node = parent_node.parent # if it reaches the root, means there is no predecessor if not parent_node: return NodeBST(None, None) succ_node = parent_node return succ_node def print_in_order(self, source=None, descending=False): """ Print the key-value pairs in the ascending/descending order of keys. Args: source: used for recursion, not for users descending: True for descending order """ self.__check_root() if not source: source = self.root if not descending: if source.left_child and source.left_child.key: self.print_in_order(source.left_child) print("{} : {}".format(source.key, source.value)) if source.right_child and source.right_child.key: self.print_in_order(source.right_child) else: if source.right_child and source.right_child.key: self.print_in_order(source.right_child) print("{} : {}".format(source.key, source.value)) if source.left_child and source.left_child.key: self.print_in_order(source.left_child) def delete(self, key): """ Delete a node with the given key. Args: key: the key of the node to be deleted """ self.__check_root() parent_node, search_node = self.__compare(key, method='search') self.__check_node(search_node) # Case 1: the node has no children nodes if (not search_node.left_child) and (not search_node.right_child): search_node.reset() # if the node has parent, delete the child of the parent if parent_node and parent_node.key: if key <= parent_node.key: parent_node.left_child = None else: parent_node.right_child = None # Case 2: the node has only one child node elif bool(search_node.left_child) != bool(search_node.right_child): child = search_node.left_child if search_node.left_child else search_node.right_child if key <= parent_node.key: parent_node.left_child = child else: parent_node.right_child = child child.parent = parent_node search_node.reset() # Case 3: the node has two children nodes else: pred = self.get_predecessor(key) self.__swap_kv(search_node, pred) if pred.left_child and pred.left_child.key: pred.parent.right_child = pred.left_child else: pred.parent.right_child = None pred.reset() # tree = BinarySearchTree() # random.seed(1) # # key_list = [random.randint(1000, 2000) for i in range(10)] # value_list = [random.randint(60, 100) for i in range(10)] # # for i in range(10): # tree.insert(key_list[i], value_list[i]) # # print('------------------------------------------') # sorted_list = sorted(key_list) # print('max: %s' % max(key_list)) # print('min: %s' % min(key_list)) # print('sorted array: {}'.format(sorted_list)) # print('------------------------------------------') # # print("min key: %s" % tree.min.key) # print("max key: %s" % tree.max.key) # print("successor of 1261 is: %s" % tree.get_successor(1261).key) # # print('----->') # tree.print_in_order() # print('----->') # # print('delete value 1867') # tree.delete(1867) # tree.search(1779, print_path=True) # # print('rotate around 1582') # tree.rotation(1582) # tree.search(1507, print_path=True)