taller = False
LH = 1
EH = 0
RH = -1


class BiTree(object):

    def __init__(self, data):
        self.data = data
        self.lChild = None
        self.rChild = None
        self.bf = 0


class AVLTree(object):

    def __init__(self):
        self.root = None

    def createAVL(self, inputList):
        if not inputList:
            return None
        for v in inputList:
            self.root = self.insertAVL(self.root, v)
        return self.root

    def leftRotate(self, parent):
        rchild = parent.rChild
        if not rchild:
            return
        parent.rChild = rchild.lChild
        rchild.lChild = parent
        return rchild

    def rightRotate(self, parent):
        lchild = parent.lChild
        if not lchild:
            return
        parent.lChild = lchild.rChild
        lchild.rChild = parent
        return lchild

    def leftBalance(self, parent):
        left = parent.lChild
        # LL情况，L节点变为此子树根节点
        # L的右孩子为原根节点
        # 原根节点的左孩子为L的右孩子
        if LH == left.bf:
            parent.bf = EH
            left.bf = EH
            parent = self.rightRotate(parent)
        # LR情况，LR节点变为此子树根节点
        # LR的左孩子为L节点，LR的右孩子为原根节点
        # L的右孩子为LR的左孩子，原根节点的左孩子为LR的右孩子
        elif RH == left.bf:
            lr = left.rChild
            if EH == lr.bf:
                parent.bf = EH
                left.bf = EH
            elif LH == lr.bf:
                parent.bf = RH
                left.bf = EH
            elif RH == lr.bf:
                parent.bf = EH
                left.bf = LH
            lr.bf = EH
            parent.lChild = self.leftRotate(left)
            parent = self.rightRotate(parent)
        return parent

    def rightBalance(self, parent):
        right = parent.rChild
        # RR情况
        if RH == right.bf:
            parent.bf = EH
            right.bf = EH
            parent = self.leftRotate(parent)
        # RL情况
        elif LH == right.bf:
            rl = right.lChild
            if EH == rl.bf:
                parent.bf = EH
                right.bf = EH
            elif LH == rl.bf:
                parent.bf = EH
                right.bf = RH
            elif RH == rl.bf:
                parent.bf = LH
                right.bf = EH
            rl.bf = EH
            parent.rChild = self.rightRotate(right)
            parent = self.leftRotate(parent)
        return parent

    def insertAVL(self, parent, v):
        global taller
        if not parent:
            parent = BiTree(v)
            taller = True
            return parent
        if v < parent.data:
            # 递归调用，必然从最小非平衡子树开始调整
            parent.lChild = self.insertAVL(parent.lChild, v)
            if taller:
                # 插之前左偏，插入左子树后失衡，执行左平衡操作
                # 平衡后无需判断上级父节点平衡情况，taller置为False
                if LH == parent.bf:
                    parent = self.leftBalance(parent)
                    taller = False
                # 插之前平衡，插入左子树后左偏
                # 此子树高度变高，返回后需判断上级父节点平衡情况，taller置为True
                elif EH == parent.bf:
                    parent.bf += 1
                    taller = True
                # 插之前右偏，插入左子树后平衡
                # 此子树高度不变，返回后无需判断上级父节点平衡情况，taller置为False
                elif RH == parent.bf:
                    parent.bf += 1
                    taller = False
        else:
            parent.rChild = self.insertAVL(parent.rChild, v)
            if taller:
                # 插之前左偏，右插后将此子树置为平衡
                if LH == parent.bf:
                    parent.bf -= 1
                    taller = False
                # 插之前平衡，右插后将此子树置为右偏
                elif EH == parent.bf:
                    parent.bf -= 1
                    taller = True
                # 插之前右偏，插入右子树后失衡，执行右平衡操作
                elif RH == parent.bf:
                    parent = self.rightBalance(parent)
                    taller = False
        return parent

    # 中序遍历进行验证
    def midTraverse(self, root):
        if not root:
            return None
        self.midTraverse(root.lChild)
        print(root.data, end=' ')
        self.midTraverse(root.rChild)


if __name__ == '__main__':
    T = AVLTree()
    l = [3, 1, 2]
    root = T.createAVL(l)
    T.midTraverse(root)