平衡二叉树的遍历/删除/新增/维护平衡因子

class Tree(object):
    '''
    树节点
    '''
    def __init__(self, data, parent, left=None, right=None, factor=0):
        self.data = data
        self.parent = parent
        self.left = left
        self.right = right
        self.factor = factor
        
class AVLTree(object):
    def __init__(self):
        self.root = None
        
    
    def insertNode(self, key):
        '''
        找到要插入的元素的位置，返回保存记录从根节点到插入节点的位置
        '''
        if self.root is None:
            self.root = Tree(key, None)
            return [self.root]
        traverse_path = []
        node = self.root
        while(node is not None):
            traverse_path.append(node)
            if key > node.data:
                if node.right is None:
                    node.right = Tree(key, node)
                    return traverse_path
                else:
                    node = node.right
            elif key < node.data:
                if node.left is None:
                    node.left = Tree(key, node)
                    return traverse_path
                else:
                    node=node.left
            else:
                print("%d is in BinaryTree" % key)
                return traverse_path
    
    def searchKeyPath(self, key):
        '''
        从根节点出发查找关键字key经过的path
        '''
        if self.root is None:
            print('root is empty')
            return 
        parent_path = []
        node = self.root
        while node is not None:
            if key < node.data:
                parent_path.append(node)
                node = node.left
            elif key > node.data:
                parent_path.append(node)
                node = node.right
            else:
                return parent_path
            
    
    def calculate_depth(self, node):
        '''
        计算节点的高度
        '''
        if (node.left is None) and (node.right is None):
            return 0
        elif node.left is None:
            height = self.calculate_depth(node.right) + 1
#             print("height right is %d"%height)
        elif node.right is None:
            height = self.calculate_depth(node.left) + 1
#             print("height left is %d"%height)
        else:
            height = max(self.calculate_depth(node.left) , self.calculate_depth(node.right)) + 1
#             print("height is %d"%height)
        
        return height
            
    def calculate_factor(self, node):
        '''
        计算节点的平衡因子
        '''
        if node.left is None and node.right is None:
            return 0
        elif node.left is None:
            return 0 - (self.calculate_depth(node.right)+1)
        elif node.right is None:
            return (self.calculate_depth(node.left)+1) - 0
        else:
            return (self.calculate_depth(node.left)+1) - (self.calculate_depth(node.right)+1)


    '''
    维持平衡因子，与网上的LL,LR,RR,RL旋转不同，这里分为两类左旋和右旋，情况差不多吧，但是树节点的摆放不一样，前/中/后序遍历是一样的
    '''
    def rotate_left(self, node):
        '''
        找到node左子树最大的节点
        '''
        # node节点是打破平衡因子的节点，需要调整
        rotate_left_path = []
        parent = node
        rotate_left_path.append(parent)
        temp = node.left
        ## 类似于LR
        if temp.right:
            parent = temp
            rotate_left_path.append(parent)
            temp = temp.right
            while temp.right:
                parent = temp
                rotate_left_path.append(parent)
                temp = temp.right
            # node 有parent
            if node.parent:
                # 设置node的parent指向
                if node.parent.data < temp.data:
                    node.parent.right = temp
                else:
                    node.parent.left = temp
                # 设置parent
                temp.parent = node.parent
                node.parent = temp
            else:
                node.parent = temp
                temp.parent = None
                self.root = temp
            # 进行右旋转    
            parent.right = temp.left
            if temp.left:
                temp.left.parent = parent
            temp.left = node.left
            node.left.parent = temp
            temp.right = node
            node.left = None
            
        # node的孩子节点没有右子树 ，类似LL
        else:
            # 如果node有父节点
            if node.parent:
                # 父节点指向新的节点即temp
                if node.parent.data < temp.data:
                    node.parent.right = temp
                else:
                    node.parent.left = temp
                    
                # 设置node 和 temp的父节点
                temp.parent = node.parent
                node.parent = temp
            # 如果node没有父节点，即node是根节点
            else:
                node.parent = temp
                temp.parent = None
                self.root = temp
                
            temp.right = node
            node.left = None
        
        # 从根节点出发找到temp节点的路径
#         rotate_left_path = self.searchKeyPath(node.data)
        rotate_left_path.extend([node])
        print('rotate node parent path ')
        for element in rotate_left_path:
            print(element.data, end=' ')
        rotate_left_path = rotate_left_path[::-1]
        for element in rotate_left_path:
            print(element.data, end=' ')
        print('end rotate node parent path')
        return rotate_left_path
        
    
    def rotate_right(self, node):
        '''
        找到node右子树最小的节点
        '''
        rotate_right_path = []
        parent = node
        rotate_right_path.append(parent)
        temp = node.right
        if temp.left:
            parent = temp
            rotate_right_path.append(parent)
            temp = temp.left
            while temp.left:
                parent = temp
                rotate_right_path.append(parent)
                temp = temp.left
            if node.parent:
                if node.parent.data > temp.data:
                    node.parent.left = temp
                else:
                    node.parent.right = temp
                temp.parent = node.parent
                node.parent = temp
            else:
                node.parent = temp
                temp.parent = None
                self.root = temp
            # 对右子树进行左旋转
            temp.left = node
            parent.left = temp.right
            if temp.right:
                temp.right.parent = parent
            temp.right = node.right
            node.right.parent = temp
            node.right = None
        else:
            # 对右子树进行左旋转
            if node.parent:
                if node.parent.data > temp.data:
                    node.parent.left = temp
                else:
                    node.parent.right = temp
                temp.parent = node.parent
                node.parent = temp
                temp.left = node
                node.right = None
            # 果然node节点没有父节点
            else:
                # 设置parent
                node.parent = temp
                temp.parent = None
                # 设置child
                self.root = temp
                temp.left = node
                node.right = None
        
        # 从根节点出发找到temp节点的路径
#         rotate_right_path = self.searchKeyPath(node.data)
        rotate_right_path.extend([node])
        print('rotate node parent path ')
        for element in rotate_right_path:
            print(element.data, end=' ')
        rotate_right_path = rotate_right_path[::-1]
        for element in rotate_right_path:
            print(element.data, end=' ')
        print('end rotate node parent path')
        return rotate_right_path
        
            
                
    
    
    def beforeAndFactorIterator(self):
        ''' before traversing
        '''
        node = self.root
        stack = []
        while(node or stack):
            while node:
                yield (node.data, node.factor)
                stack.append(node)
                node = node.left
            node = stack.pop()
            node = node.right
            
    def beforeIterator(self):
        ''' before traversing
        '''
        node = self.root
        stack = []
        while(node or stack):
            while node:
                yield node.data
                stack.append(node)
                node = node.left
            node = stack.pop()
            node = node.right
            
    def __iter__(self):
        ''' inorder traversing
        '''
        node = self.root
        stack = []
        while (node or stack):
            while node :
                stack.append(node)
                node = node.left
            node = stack.pop()
            yield node.data
            node = node.right
            

    
    def postIterator(self):
        ''' postorder traversing
        '''
        node = self.root
        stack = []
        while node or stack:
            while node and node.traversed_right == 0 :
                node.traversed_right = 1 # 表示已经入list，list中的节点不能再向左访问
                stack.append(node)
                node = node.left
            
            node = stack.pop()
            if node.right and node.traversed_right != 2:
                node.traversed_right = 2
                stack.append(node)
                node = node.right
            else:
                yield node.data
                if len(stack) == 0:
                    break  
  
  
if __name__ == '__main__':
#     lis = [62, 58, 88, 48, 73, 99, 35, 51, 93, 29, 37, 49, 56, 36, 50]
#     lis = [20, 10, 30, 15, 5, 1]
    
    # test accuracy
    lis = []
    for element in range(1,100,1):
        lis.append(element)
    print(lis)    
    avl_tree = AVLTree()
    for i in range(len(lis)):
        traverse_path = avl_tree.insertNode(lis[i])
        # 倒序遍历
        while traverse_path :
            node = traverse_path.pop()
            factor = avl_tree.calculate_factor(node)
            if factor >= 2 :
                for i in avl_tree.beforeIterator():
                    print(i, end=" ")
                    
                print('end')    
                print("rotate_left %d" % node.data)
                rotate_left_path = avl_tree.rotate_left(node)
                factor = avl_tree.calculate_factor(node)
                for element in rotate_left_path:
                    if element not in traverse_path:
                        traverse_path.append(element)
            elif factor <= -2 :
                for i in avl_tree.beforeIterator():
                    print(i, end=" ")
                    
                print('before rotate_right')    
                print("rotate_right %d" % node.data)
                rotate_right_path = avl_tree.rotate_right(node)
                factor = avl_tree.calculate_factor(node)
                for element in rotate_right_path:
                    if element not in traverse_path:
                        traverse_path.append(element)
            
            node.factor = factor
            for i, factor in avl_tree.beforeAndFactorIterator():
                print("d=%d f=%d"%(i, factor), end=", ")
            print("end for")
            
            
    for node in avl_tree:
        print(node, end=' ')
    print('end
 before traversing')
    
    for i, factor in avl_tree.beforeAndFactorIterator():
        print("d=%d f=%d"%(i, factor), end=", ")
    print('end')

LL/LR/RR/RL旋转方式见链接

Reference

[1] https://www.cnblogs.com/sfencs-hcy/p/10356467.html