应该对此代码进行哪些更改才能正确识别平衡树？

class TreeNode:
    def __init__(self, key):
        self.key = key
        self.left = None
        self.right = None
        # 对于AVL树，添加平衡因子
        self.height = 1
        # 对于红黑树，添加颜色属性
        # self.color = 'RED' or 'BLACK'

def insert(node, key):
    if not node:
        return TreeNode(key)
    if key < node.key:
        node.left = insert(node.left, key)
    else:
        node.right = insert(node.right, key)

    node.height = 1 + max(get_height(node.left), get_height(node.right))
    balance = get_balance(node)

    # 左左情况
    if balance > 1 and key < node.left.key:
        return right_rotate(node)
    # 右右情况
    if balance < -1 and key > node.right.key:
        return left_rotate(node)
    # 左右情况
    if balance > 1 and key > node.left.key:
        node.left = left_rotate(node.left)
        return right_rotate(node)
    # 右左情况
    if balance < -1 and key < node.right.key:
        node.right = right_rotate(node.right)
        return left_rotate(node)

    return node

def left_rotate(z):
    y = z.right
    T2 = y.left

    y.left = z
    z.right = T2

    z.height = 1 + max(get_height(z.left), get_height(z.right))
    y.height = 1 + max(get_height(y.left), get_height(y.right))

    return y

def right_rotate(z):
    y = z.left
    T3 = y.right

    y.right = z
    z.left = T3

    z.height = 1 + max(get_height(z.left), get_height(z.right))
    y.height = 1 + max(get_height(y.left), get_height(y.right))

    return y

def get_height(node):
    if not node:
        return 0
    return node.height

def get_balance(node):
    if not node:
        return 0
    return get_height(node.left) - get_height(node.right)

def pre_order(node):
    if not node:
        return
    print("{0} ".format(node.key), end="")
    pre_order(node.left)
    pre_order(node.right)

root = None
keys = [10, 20, 30, 40, 50, 25]

for key in keys:
    root = insert(root, key)

print("Preorder traversal of the constructed AVL tree is")
pre_order(root)
print()