python实现二分搜索树以及改进后的AVL Tree

二分搜索树讲解
python代码：

class TreeNode(object):    def __init__(self, key, val, parent=None, left_child=None, right_child=None):        self.key = key         self.val = val         self.parent = parent # 指向该节点的双亲节点        self.left_child = left_child         self.right_child = right_child    def is_left_child(self):        return self.parent and self.parent.left_child == self    def is_right_child(self):        return self.parent and self.parent.right_child == self    def is_root(self):        return not self.parent    def is_leaf(self):        return not (self.left_child or self.right_child)    def has_any_child(self):        return self.left_child or self.right_child    def has_both_childs(self):        return self.left_child and self.right_child    def find_successor(self):        """        寻找该节点的successor        successor必须是leaf node 或 只有一个子节点        """        if self.right_child:            return self.right_child.find_min()        elif self.is_left_child():            return self.parent        else:            self.parent.right_child = None            succ = self.parent.find_successor()            self.parent.right_child = self            return succ    def find_min(self):        # 寻找该节点及其子节点中key最小的节点        current_node = self        while current_node.left_child:            current_node = current_node.left_child        return current_node    def splice(self):        """        去除successor，并将其子节点和其双亲节点嫁接        """        if self.is_leaf():            if self.is_left_child():                self.parent.left_child = None            else:                self.parent.right_child = None        elif self.has_any_child():            if self.left_child:                if self.is_left_child():                    self.parent.left_child = self.left_child                else:                    self.parent.right_child = self.left_child                self.left_child.parent = self.parent            else:                if self.is_left_child():                    self.parent.left_child = self.right_child                else:                    self.parent.right_child = self.right_child                self.right_child.parent = self.parent    def replace_node_data_childs(self, new_key, new_val, new_left_child, new_right_child):        self.key = new_key        self.val = new_val        self.left_child = new_left_child        self.right_child = new_right_child    def __iter__(self):        if self:            if self.left_child:                for elem in self.left_child:                    yield elem            yield self.key            if self.right_child:                for elem in self.right_child:                    yield elemclass BinSearchTree(object):    def __init__(self, root=None):        self.root = None        self._size = 0    def __len__(self):        return self._size    def __setitem__(self, k, v):        if self.root:            self.__put(k, v, self.root)        else:  # the search tree is empty            self.root = TreeNode(k, v)        self._size += 1    def __put(self, k, v, current_node):        if k < current_node.key:            if current_node.left_child:                self.__put(k, v, current_node.left_child)            else:                current_node.left_child = TreeNode(k, v, parent=current_node)        elif k > current_node.key:            if current_node.right_child:                self.__put(k, v, current_node.right_child)            else:                current_node.right_child = TreeNode(k, v, parent=current_node)        else:  # the k == current_node.key            current_node.val = v    def __getitem__(self, k):        if self._size > 1:            item = self.__get(k, self.root)            if item:                return self.__get(k, self.root).val            else:                return None        elif self._size == 1 and self.root.key == k:            return self.root.val        else:            return None            # if self.root:            #   if self._size == 1: # the tree only has root node            #       if self.root.key == k:            #           return self.root.val            #       else:            #           return None            #   else:            #       if self.__get(k, self.root):            #           return self.__get(k, self.root).val            #       return None            # else: # the tree is empty            #   return None    def __get(self, k, current_node):        if not current_node:            return None        elif k == current_node.key:            return current_node        elif k < current_node.key:            return self.__get(k, current_node.left_child)        else:            return self.__get(k, current_node.right_child)            # if current_node:            #   if k < current_node.key:            #       return self.__get(k, current_node.left_child)            #   elif k > current_node.key:            #       return self.__get(k, current_node.right_child)            #   else:            #       return current_node            # else: # the tree doesn't has the key            #   return None    def __contains__(self, k):        if self.__getitem__(k):            return True        else:            return False    def __delitem__(self, k):        if self._size > 1:            del_node = self.__get(k, self.root)            if del_node:                self.__remove(del_node)                self._size -= 1            else:                raise KeyError('Error, key not in tree')        elif self._size == 1 and self.root.val == k:            self.root = None            self._size -= 1        else:            raise KeyError('Error, key not in tree')    def __remove(self, del_node):        # if the node that would be deleted has no childs        if del_node.is_leaf():            if del_node.is_left_child():                del_node.parent.left_child = None            else:                del_node.parent.right_child = None        elif del_node.has_both_childs():            successor = del_node.find_successor()            successor.splice()            del_node.key = successor.key            del_node.val = successor.val        else:  # the node that would be deleted has a child            if del_node.left_child:                if del_node.is_left_child():                    del_node.parent.left_child, del_node.left_child.parent = \                        del_node.left_child, del_node.parent                elif del_node.is_right_child():                    del_node.parent.right_child, del_node.left_child.parent = \                        del_node.left_child, del_node.parent                else:  # the node that would be deleted is root node                    del_node.replace_node_data_childs(del_node.left_child.key,                                                      del_node.left_child.val,                                                      del_node.left_child.left_child,                                                      del_node.left_child.right_child)            else:  # the node that would be deleted has right child                if del_node.is_left_child():                    del_node.parent.left_child, del_node.right_child.parent = \                        del_node.right_child, del_node.parent                elif del_node.is_right_child():                    del_node.parent.right_child, del_node.right_child.parent = \                        del_node.right_child, del_node.parent                else:                    del_node.replace_node_data_childs(del_node.right_child.key,                                                      del_node.right_child.val,                                                      del_node.right_child.left_child,                                                      del_node.right_child.right_child)    def __iter__(self):        if self.root:            for k in self.root:                yield k        else:            raise KeyError('the tree is empty')if __name__ == '__main__':    mytree = BinSearchTree()    # setitem    mytree[3] = "red"    mytree[4] = "blue"    mytree[6] = "yellow"    mytree[2] = "at"    mytree[5] = "green"    # delete item    del mytree[4]    for k in mytree:        print(k)

AVL树讲解
python代码：

from bin_search_tree import BinSearchTree, TreeNodeclass AVLTreeNode(TreeNode):    def __init__(self, key, val, parent=None, left_child=None, right_child=None):        super().__init__(key, val, parent, left_child, right_child)        self.balance_factor = 0class AVLTree(BinSearchTree):    def __init__(self, root=None):        super().__init__(root)    def __setitem__(self, k, v):        """        override the super().__setitem__        """        if self.root:            self.__put(k, v, self.root)        else:  # the search tree is empty            self.root = AVLTreeNode(k, v)        self._size += 1    def __put(self, k, v, current_node):        """        override the super().__put        """        if k < current_node.key:            if current_node.left_child:                self.__put(k, v, current_node.left_child)            else:                current_node.left_child = AVLTreeNode(k, v, parent=current_node)                self.__update_balance(current_node.left_child)        elif k > current_node.key:            if current_node.right_child:                self.__put(k, v, current_node.right_child)            else:                current_node.right_child = AVLTreeNode(k, v, parent=current_node)                self.__update_balance(current_node.right_child)        else:            current_node.val = v    def __update_balance(self, current_node):        if current_node.balance_factor < -1 or \           current_node.balance_factor > 1:           self.__rebalance(current_node)           return        if current_node.parent != None:            if current_node.is_left_child():                current_node.parent.balance_factor += 1            elif current_node.is_right_child():                current_node.parent.balance_factor -= 1            if current_node.parent.balance_factor != 0:                self.__update_balance(current_node.parent)    def __rebalance(self, current_node):        if current_node.balance_factor < -1: # need left rotation            if current_node.right_child.balance_factor > 0:                    self.__right_rotate(current_node.right_child)            self.__left_rotate(current_node)        elif current_node.balance_factor > 1: # need right rotation            if current_node.left_child.balance_factor < 0:                self.__left_rotate(current_node.left_child)            self.__right_rotate(current_node)    def __left_rotate(self, rot_root):        """        the node's balance factor < -1        """        # 将rot root的右子节点设为new root        new_root = rot_root.right_child        # 将new root的左子节点设置为rot root的右子节点        rot_root.right_child = new_root.left_child        if new_root.left_child:            new_root.left_child.parent = rot_root        # 将new root的双亲节点设置为rot root的双亲节点        new_root.parent = rot_root.parent        if rot_root.is_root():            self.root = new_root        else:            if rot_root.is_left_child():                rot_root.parent.left_child = new_root            else:                rot_root.parent.right_child = new_root        # 将rot root的双亲节点设置为new root(rot root为new root的左子节点)        new_root.left_child = rot_root        rot_root.parent = new_root        # critical steps        # update the balance factor both of rot root and new root        rot_root.balance_factor = rot_root.balance_factor + 1 - min(new_root.balance_factor, 0)        new_root.balance_factor = new_root.balance_factor + 1 + max(rot_root.balance_factor, 0)    def __right_rotate(self, rot_root):        """        the node's balance factor > 1        """        # 将rot root的左子节点设为new root        new_root = rot_root.left_child        # 将new root的右子节点设置为rot root的左子节点        rot_root.left_child = new_root.right_child        if new_root.right_child:            new_root.right_child.parent = rot_root        # 将new root的双亲节点设置为rot root的双亲节点        new_root.parent = rot_root.parent        if rot_root.is_root():            self.root = new_root        else:            if rot_root.is_left_child():                rot_root.parent.left_child = new_root            else:                rot_root.parent.right_child = new_root        # 将rot root的双亲节点设置为new root(rot root为new root的右子节点)        new_root.right_child = rot_root        rot_root.parent = new_root        # critical steps        # update the balance factor both of rot root and new root        rot_root.balance_factor = rot_root.balance_factor - 1 + min(new_root.balance_factor, 0)        new_root.balance_factor = new_root.balance_factor - 1 - max(rot_root.balance_factor, 0)if __name__ == '__main__':    mytree = AVLTree()    # setitem    mytree[3] = "red"    mytree[4] = "blue"    mytree[6] = "yellow"    # mytree[2] = "at"    # mytree[5] = "green"    print('root:%s|balance_factor:%s' % (mytree.root.key, mytree.root.balance_factor))    print('root.left_child:%s' % mytree.root.left_child.key)    print('root.right_child:%s' % mytree.root.right_child.key)