【模板】AVL树

发表于 分类于 阅读次数： Valine：

AVL树的模板（仅插入）

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
#include<iostream>

#define L 0
#define R 1
using namespace std;
struct node {
    int num, h;
    node *l, *r;
};

node *init(int num) {
    node *cur = new node;
    cur->num = num;
    cur->h = 1;
    cur->l = cur->r = NULL;
    return cur;
}

int height(node *cur) {
    if (!cur) return 0;
    else return cur->h;
}

int upd(node *cur) {
    return max(height(cur->l), height(cur->r)) + 1;
}

node *rt(node *cur, bool type) {
    node *tmp;
    if (type == R) {
        tmp = cur->l;
        cur->l = cur->l->r;
        tmp->r = cur;
    } else {
        tmp = cur->r;
        cur->r = cur->r->l;
        tmp->l = cur;
    }
    cur->h = upd(cur);
    tmp->h = upd(tmp);
    return tmp;
}

int bal(node *cur) {
    return height(cur->l) - height(cur->r);
}

node *modi(node *cur) {
    if (bal(cur) == 2) {
        if (bal(cur->l) == 1) cur = rt(cur, R);//LL
        else if (bal(cur->l) == -1) { // LR
            cur->l = rt(cur->l, L);
            cur = rt(cur, R);
        }
    } else if (bal(cur) == -2) {
        if (bal(cur->r) == -1) cur = rt(cur, L);//RR
        else if (bal(cur->r) == 1) {//RL
            cur->r = rt(cur->r, R);
            cur = rt(cur, L);
        }
    }
    return cur;
}

node *insert(node *cur, int num) {
    if (!cur) return init(num);
    if (num < cur->num) cur->l = insert(cur->l, num);
    else cur->r = insert(cur->r, num);
    cur->h = upd(cur);
    return modi(cur);
}

void dfs(node *cur) {
    if (!cur) return;
    printf("%d,", cur->num);
    dfs(cur->l);
    dfs(cur->r);
    delete cur;
}

int main() {
    int num;
    scanf("%d,", &num);
    node *root = init(num);
    while (~scanf("%d,", &num)) root = insert(root, num);
    dfs(root);
    return 0;
}