Python实现自平衡二叉树AVL

# -*- coding: utf-8 -*-
from enum import Enum

#参考http://blog.csdn.net/niteip/article/details/11840691/
#参考https://www.cnblogs.com/suimeng/p/4560056.html
#todo 还没有考虑高度的增减，只考虑了平衡因子
#todo 加上非递归遍历二叉树

class BFStatus(Enum):
    # 左子树的高度减去右子树的高度 RH,EH,LH分别表示右子树较高，左右子树等高，左子树较高
    LH = 1
    EH = 0
    RH = -1

class TreeNode(object):
    def __init__(self):
        self.lson = None
        self.rson = None
        self.pnode = None
        self.data = None
        self.freq = 0
        self.hgt = -1
        self.bfstatus = BFStatus().EH

class AVLTree(object):
    def __init__(self,root_node):
        self.root = TreeNode

    def height(self,node):
        if node is not None:
            return node.hgt
        return -1

    def max(self,campa,campb):
        if campa > campb:
            return campa
        else:
            return campb

    def single_rotate_right(self, node):  # 右单旋转 RR
        parent_node = node.pnode
        node1 = node.lson
        node.lson = node1.rson
        node1.rson = node
        if parent_node.lson == node:
            parent_node.lson = node1
        else:
            parent_node.rson = node1
        node1.pnode = parent_node
        node.pnode = node1

    def single_rotate_left(self, node):  # 左单旋转 LL
        parent_node = node.pnode
        node1 = node.rson
        node.rson = node1.lson
        node1.lson = node
        if parent_node.lson == node:
            parent_node.lson = node1
        else:
            parent_node.rson = node1
        node1.pnode = parent_node
        node.pnode = node1

    def LeftBalance(self,node):
        lchild = node.lson
        BFStatus_code = BFStatus()
        if lchild.bfstatus == BFStatus_code.EH:
            node.bfstatus = lchild.bfstatus = BFStatus_code.LH
            lchild.bfstatus = BFStatus.RH
            self.single_rotate_right(node)
        elif lchild.bfstatus == BFStatus_code.LH:#LL型
            node.bfstatus = lchild.bfstatus = BFStatus.EH
            self.single_rotate_right(node)
        elif lchild.bfstatus == BFStatus.RH:#LR型
            rchild = lchild.rson
            if rchild.bfstatus == BFStatus.EH:
                node.bfstatus = BFStatus.RH
                lchild.bfstatus = BFStatus.EH
            elif rchild.bfstatus == BFStatus.LH:
                node.bfstatus = BFStatus.EH
                lchild.bfstatus = BFStatus.LH
            elif rchild.bfstatus == BFStatus.RH:
                node.bfstatus = BFStatus.EH
                lchild.bfstatus = BFStatus.EH
            rchild.bfstatus = BFStatus.EH
            self.single_rotate_left(lchild)
            self.single_rotate_right(node)

    def RightBalance(self,node):
        rchild = node.rson
        BFStatus_code = BFStatus()
        if rchild.bfstatus == BFStatus_code.RH:
            node.bfstatus = node.bfstatus = BFStatus_code.EH
            self.single_rotate_left(node)
        elif rchild.bfstatus == BFStatus_code.EH:#LL型
            node.bfstatus = rchild.bfstatus = BFStatus.RH
            rchild.bfstatus = BFStatus.LH
            self.single_rotate_left(node)
        elif rchild.bfstatus == BFStatus.LH:#LR型
            lchild = rchild.lson
            if lchild.bfstatus == BFStatus.LH:
                node.bfstatus = BFStatus.EH
                rchild.bfstatus = BFStatus.RH
            elif lchild.bfstatus == BFStatus.RH:
                node.bfstatus = BFStatus.LH
                rchild.bfstatus = BFStatus.EH
            elif lchild.bfstatus == BFStatus.EH:
                node.bfstatus = BFStatus.EH
                rchild.bfstatus = BFStatus.EH
            lchild.bfstatus = BFStatus.EH
            self.single_rotate_right(rchild)
            self.single_rotate_left(node)

    def insertpri(self,node,data,stack_list,taller_param=True):
        if node == None:
            node = TreeNode()
            node.data = data
            node.pnode = stack_list[-1]
        else:
            stack_list.append(node)
            if node.data < data:
                (sign,taller) = self.insertpri(node.lson,data,stack_list)
                if sign == False:
                    return (False,taller_param)

                if taller == True:

                    if node.bfstatus == BFStatus().LH:
                        self.LeftBalance(node)
                        taller_param = False
                    elif node.bfstatus == BFStatus().EH:
                        node.bfstatus =  BFStatus().LH
                        taller_param = True
                    elif node.bfstatus == BFStatus().RH:
                        node.bfstatus =  BFStatus().EH
                        taller_param = False

            elif node.data > data:
                stack_list.append(node)
                if node.data < data:
                    (sign, taller) = self.insertpri(node.rson, data, stack_list)
                    if sign == False:
                        return (False, taller_param)

                    if taller == True:
                        if node.bfstatus == BFStatus().LH:
                            node.bfstatus = BFStatus().EH
                            taller_param = False
                        elif node.bfstatus == BFStatus().EH:
                            node.bfstatus = BFStatus().RH
                            taller_param = True
                        elif node.bfstatus == BFStatus().RH:
                            self.RightBalance(node)
                            taller_param = False
            else:
                node.freq += 1

        return (True,taller_param)

    def insert(self,data):
        stack_list = []
        self.insertpri(self.root,data,stack_list)

    def searchpri(self,node,data):
        if node is None:
            return None

        elif node.data > data:
            self.searchpri(node.lson,data)

        elif node.data < data:
            self.searchpri(node.rson,data)

        else:
            return node

    def search(self,data):
        self.searchpri(self.root,data)

    def go_far_left(self,node):
        if node.lson is None:
            return node
        else:
            self.go_far_left(node.lson)

    def go_far_right(self,node):
        if node.rson is None:
            return node
        else:
            self.go_far_right(node.rson)

    def delete_left(self,node,bfchild):
        #当bf为-1或1变为0，或者孩子为空时说明子树高降低
        if node.lson is None or (bfchild != BFStatus().EH and node.lson.bfstatus == BFStatus().EH):
            #左子树树高降低
            if node.bfstatus == BFStatus().EH:
                node.bfstatus = BFStatus().RH
            elif node.bfstatus == BFStatus().LH:
                node.bfstatus = BFStatus().EH
            elif node.bfstatus == BFStatus().RH:#原本右子树比较高
                self.RightBalance(node)

    def delete_right(self,node,bfchild):
        #当bf为LH或RH变为EH，或者孩子为空时说明子树高降低
        if node.rson is None or (bfchild != BFStatus().EH and node.rson.bfstatus == BFStatus().EH):
            #左子树树高降低
            if node.bfstatus == BFStatus().EH:
                node.bfstatus = BFStatus().LH
            elif node.bfstatus == BFStatus().RH:
                node.bfstatus = BFStatus().EH
            elif node.bfstatus == BFStatus().LH:#原本左子树比较高
                self.LeftBalance(node)

    def deletepri(self,node,data):
        bfchild = BFStatus().EH
        if node == None:
            return None
        if node.data > data:
            bfchild = node.lson.bfstatus
            node.lson = self.deletepri(node.lson,data)
            self.delete_left(node,bfchild)

        elif node.data < data:
            bfchild = node.rson.bfstatus
            node.rson = self.deletepri(node.rson, data)
            self.delete_left(node, bfchild)

        else:
            if node.lson is not None:   #不是叶子结点并且具有直接前驱
                far_right_node = self.go_far_right(node.lson)#直接前驱
                node.data = far_right_node.data
                node.lson = self.deletepri(node.lson,far_right_node.data)#可以确定，删除的节点为当前节点的左子树的某一个节点
                self.delete_left(node,bfchild)
            elif node.rson is not None:
                far_left_node = self.go_far_left(node.rson)
                node.data = far_left_node.data
                node.rson = self.deletepri(node.rson,far_left_node.data)
                self.delete_right(node,bfchild)
            else:#叶子结点
                node = None

        return node

    def delete(self,data):
        self.deletepri(self.root,data)

    def insubtree(self,node):
        if node is None:
            return
        self.insubtree(node.lson)
        print(node.data)
        self.insubtree(node.rson)

    def traversal(self):
        self.insubtree(self.root)

if __name__ == '__main__':
    root_node = TreeNode()
    tree_obj = AVLTree(root_node)