AVL建树实现

class TreeNode:
    """
        树节点结构，加了左右子树的高度
    """
    def __init__(self, data):
        self.data = data
        self.left = None
        self.right = None
        self.lheight = 0
        self.rheight = 0

class BalanceTree:
    # 建树，输入列表数据，返回树的根结点
    def build(self, data):
        root = TreeNode(data[0])
        data = data[1:]
        for i in data:
            root = self.add(i, root)
        self.midOrder(root)
        return root

    # 中序遍历，检验平衡二叉树
    def midOrder(self, root):
        if not root:
            return
        self.midOrder(root.left)
        print(root.data)
        self.midOrder(root.right)

    # 添加结点数据
    def add(self, data, root):
        nodeNew = TreeNode(data)
        # tmp记录插入新结点的双亲节点, p遍历结点, nodeTag需要进行调整的结点，存在不平衡现象,
        # nodeTagF需要调整的结点的双亲结点,
        # tagF记录每一次遍历的前一次结点，用于生成当前结点的双亲结点
        tmp, p, nodeTag, nodeTagF, tagF = root, root, None, None, None
        # 遍历找到插入新结点的位置
        while p:
            tmp = p
            if data < p.data:
                p.lheight += 1
                # 检验是否存在不平衡现象
                if abs(self.checkTree(p)) > 1:
                    nodeTag = p
                    nodeTagF = tagF
                # 保存下一次遍历结点的双亲结点
                tagF = p if p.left else tagF
                p = p.left
            else:
                p.rheight += 1
                if abs(self.checkTree(p)) > 1:
                    nodeTag = p
                    nodeTagF = tagF
                tagF = p if p.right else tagF
                p = p.right
        # 插入新结点
        if data < tmp.data:
            tmp.left = nodeNew
        else:
            tmp.right = nodeNew
        # 判断是否需要进行调整二叉树
        if nodeTag != None:
            # 如果当前不平衡的结点为根结点，调整后直接返回结点
            if nodeTag == root:
                return self.adjust(nodeTag)
            else:
                # 不是根结点，需要将调整后的子树重新插入，并调整双亲节点的子树高度
                if nodeTagF.left == nodeTag:
                    balanceNode = self.adjust(nodeTag)
                    nodeTagF.left = balanceNode
                    nodeTagF.lheight = max(balanceNode.rheight, balanceNode.lheight)
                else:
                    balanceNode = self.adjust(nodeTag)
                    nodeTagF.right = balanceNode
                    nodeTagF.rheight = max(balanceNode.rheight, balanceNode.lheight)
        return root

    def checkTree(self, root):
        return root.lheight-root.rheight

    def adjust(self, nodeTag):
        tag = self.checkTree(nodeTag)
        if tag == 2:
            # ll调整
            if self.checkTree(nodeTag.left) == 1:
                return self.ll(nodeTag)
            else:
                # lr调整
                return self.lr(nodeTag)
        else:
            # rr调整
            if self.checkTree(nodeTag.right) == -1:
                return self.rr(nodeTag)
            else:
                # rl调整
                return self.rl(nodeTag)

    def ll(self, node):
        nodeRoot = node.left
        nodeTmp = nodeRoot.right
        nodeRoot.right = node
        node.left = nodeTmp
        # ll结构中，最后作为根结点的结点，可能不存在右子树，所以需要特别判断
        node.lheight = 0 if not nodeTmp else max(nodeTmp.rheight, nodeTmp.lheight)
        nodeRoot.rheight = max(node.rheight, node.lheight)
        return nodeRoot

    def rr(self, node):
        nodeRoot = node.right
        nodeTmp = nodeRoot.left
        nodeRoot.left = node
        node.right = nodeTmp
        # rr结构中，最后作为根结点的结点，可能不存在左子树，所以需要特别判断
        node.rheight = 0 if not nodeTmp else max(nodeTmp.rheight, nodeTmp.lheight)
        nodeRoot.lheight = max(node.rheight, node.lheight)
        return nodeRoot

    def lr(self, node):
        nodeRoot = node.left.right
        nodeLtmp = nodeRoot.left
        nodeRtmp = nodeRoot.right
        nodeRoot.left = node.left
        nodeRoot.right = node
        node.left = nodeRtmp
        nodeRoot.left.right = nodeLtmp
        # lr结构中，最后作为根结点的结点，有可能只有一个子树，所以两个子树高度赋值时都需要做特别判断
        node.lheight = 0 if not nodeRtmp else max(nodeRtmp.rheight, nodeRtmp.lheight)
        nodeRoot.left.rheight = 0 if not nodeLtmp else max(nodeLtmp.rheight, nodeLtmp.lheight)
        nodeRoot.lheight = max(nodeRoot.left.lheight, nodeRoot.left.rheight)
        nodeRoot.rheight = max(nodeRoot.right.lheight, nodeRoot.right.rheight)
        return nodeRoot

    def rl(self, node):
        nodeRoot = node.right.left
        nodeLtmp = nodeRoot.left
        nodeRtmp = nodeRoot.right
        nodeRoot.right = node.right
        nodeRoot.left = node
        node.right = nodeLtmp
        nodeRoot.right.left = nodeRtmp
        # rl结构中，最后作为根结点的结点，有可能只有一个子树，所以两个子树高度赋值时都需要做特别判断
        node.rheight = 0 if not nodeLtmp else max(nodeLtmp.rheight, nodeLtmp.lheight)
        nodeRoot.right.lheight = 0 if not nodeRtmp else max(nodeRtmp.rheight, nodeRtmp.lheight)
        nodeRoot.lheight = max(nodeRoot.left.lheight, nodeRoot.left.rheight)
        nodeRoot.rheight = max(nodeRoot.right.lheight, nodeRoot.right.rheight)
        return nodeRoot

if __name__ == '__main__':
    ll = [4,2,5,6,1,9,3,8,7]
    ss = BalanceTree()
    ss.build(ll)