平衡二叉樹的遍歷/刪除/新增/維護平衡因子

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\n 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