「CodeNote」 AVL

Posted by Dawn-K's Blog on March 8, 2021

概述

AVL是一种平衡二叉树. 其本身是一种BST(二叉搜索树), 但是它可以通过旋转自身来平衡左右子树, 让左右子树的高度差不超过1. 难点在于旋转和维护高度.

代码

以下是PAT甲级1123的AC代码, 题中保证了插入树的所有元素均不相同.

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
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
#include <bits/stdc++.h>

using namespace std;

struct AVL {
    int val;
    AVL *l, *r;
    int h;

    // 切记高度初始化为1
    AVL(int val) : val(val), l(nullptr), r(nullptr), h(1) {}
};

typedef AVL *T;

// 如果每次都递归求树高的话,那么插入单个元素的复杂度是Nlogn(无法严谨证明),插入n个元素的复杂度就是N^2*logn,也就是大概1000以下的没有问题
// 如果每次都采用直接维护的h的话,插入单个元素复杂度是logn(树高),
// 采用维护树高的方式的话,只需要在左旋和右旋,以及插入后修改树高即可.要注意初始化的节点的树高为1.
int getH(T rt) {
    if (rt == nullptr) {
        return 0;
    }
    //    递归方式求树的高度
    //    int HeightInCur = max(getH(rt->l), getH(rt->r)) + 1;
    return rt->h;
}

T L(T rt) {
    if (rt == nullptr) {
        return nullptr;
    }
    T p = rt->r;
    rt->r = p->l;
    p->l = rt;
    // 注意下面两行的顺序,一定是先更新p,后更新rt,因为rt已经是p的儿子了,必须儿子先更新,才能更新父节点
    // 右旋同理
    rt->h = max(getH(rt->l), getH(rt->r)) + 1;
    p->h = max(getH(p->l), getH(p->r)) + 1;
    return p;
}

T R(T rt) {
    if (rt == nullptr) {
        return nullptr;
    }
    T p = rt->l;
    rt->l = p->r;
    p->r = rt;
    rt->h = max(getH(rt->l), getH(rt->r)) + 1;
    p->h = max(getH(p->l), getH(p->r)) + 1;
    return p;
}

T LR(T rt) {
    rt->l = L(rt->l);
    return R(rt);
}

T RL(T rt) {
    rt->r = R(rt->r);
    return L(rt);
}

T insert(T rt, int val) {
    if (rt == nullptr) {
        return new AVL(val);
    }
    if (val < rt->val) {
        rt->l = insert(rt->l, val);
        if (getH(rt->l) - getH(rt->r) == 2) {
            if (val < rt->l->val) {  // 左子树的左子树
                rt = R(rt);
            } else {  // 左子树的右子树
                rt = LR(rt);
            }
        }
    } else {
        rt->r = insert(rt->r, val);
        if (getH(rt->r) - getH(rt->l) == 2) {
            if (val < rt->r->val) {  // 右子树的左子树
                rt = RL(rt);
            } else {  // 右子树的右子树
                rt = L(rt);
            }
        }
    }
    rt->h = max(getH(rt->l), getH(rt->r)) + 1;
    return rt;
}

vector<int> level;
map<int, int> mp;

void printLevel(T rt) {
    queue<pair<T, int>> q;
    q.push(make_pair(rt, 1));
    while (!q.empty()) {
        pair<T, int> cur = q.front();
        q.pop();
        level.push_back(cur.first->val);
        mp[cur.second] = 1;
        if (cur.first->l) {
            q.push(make_pair(cur.first->l, cur.second * 2));
        }
        if (cur.first->r) {
            q.push(make_pair(cur.first->r, cur.second * 2 + 1));
        }
    }
}

int n;

// 判断是否是完全二叉树,依赖于层序遍历实现.
bool isComplete() {
    for (int i = 1; i <= n; ++i) {
        if (!mp.count(i)) {
            return 0;
        }
    }
    return 1;
}

int main() {
    cin >> n;
    T rt = nullptr;
    for (int i = 0; i < n; ++i) {
        int tmp;
        cin >> tmp;
        rt = insert(rt, tmp);
    }
    printLevel(rt);
    for (int i = 0; i < n; ++i) {
        if (i) {
            cout << " ";
        }
        cout << level[i];
    }
    cout << endl;
    if (isComplete()) {
        cout << "YES" << endl;
    } else {
        cout << "NO" << endl;
    }

    return 0;
}