PAT(Advanced Level) Practice（with python）——1066 Root of AVL Tree

Code

class Node:
    def __init__(self, key):
        self.key = key
        self.left = None
        self.right = None
        self.height = 1

class AVLTree:
    def __init__(self):
        self.root = None

    # 获取节点的高度
    def get_height(self, node):
        if node is None:
            return 0
        return node.height

    # 获取节点的平衡因子
    def get_balance_factor(self, node):
        if node is None:
            return 0
        return self.get_height(node.left) - self.get_height(node.right)

    # 更新节点的高度
    def update_height(self, node):
        if node is None:
            return
        node.height = max(self.get_height(node.left), self.get_height(node.right)) + 1

    # 执行左旋操作
    def left_rotate(self, z):
        y = z.right
        T2 = y.left

        y.left = z
        z.right = T2

        self.update_height(z)
        self.update_height(y)

        return y

    # 执行右旋操作
    def right_rotate(self, z):
        y = z.left
        T3 = y.right

        y.right = z
        z.left = T3

        self.update_height(z)
        self.update_height(y)

        return y

    # 插入节点
    def insert_node(self, root, key):
        if root is None:
            return Node(key)
        elif key < root.key:
            root.left = self.insert_node(root.left, key)
        else:
            root.right = self.insert_node(root.right, key)

        self.update_height(root)

        balance_factor = self.get_balance_factor(root)

        # 左左情况
        if balance_factor > 1 and key < root.left.key:
            return self.right_rotate(root)

        # 右右情况
        if balance_factor < -1 and key > root.right.key:
            return self.left_rotate(root)

        # 左右情况
        if balance_factor > 1 and key > root.left.key:
            root.left = self.left_rotate(root.left)
            return self.right_rotate(root)

        # 右左情况
        if balance_factor < -1 and key < root.right.key:
            root.right = self.right_rotate(root.right)
            return self.left_rotate(root)

        return root

    def insert(self, key):
        self.root = self.insert_node(self.root, key)

    # 中序遍历树
    def inorder_traversal(self, root):
        if root is not None:
            self.inorder_traversal(root.left)
            print(root.key, end=" ")
            self.inorder_traversal(root.right)

    # 获取树的高度
    def get_tree_height(self):
        return self.get_height(self.root)


# 示例代码的使用
avl_tree = AVLTree()

N= int(input())
a = map(int,input().split())
for item in a:
    avl_tree.insert(item)
print(avl_tree.root.key)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111