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
| template <typename K, typename V>
struct AVLNode {
K key;
V value;
AVLNode *left, *right;
AVLNode(K key, V value) : key(key), value(value), left(NULL), right(NULL) {}
};
template <typename K, typename V>
class AVLTree {
typedef AVLNode<K, V> Node;
Node* root;
Node* put(Node* h, K const& key, V const& value) {
if (h == nullptr) return new Node(key, value);
if (key < h->key)
h->left = put(h->left, key, value);
else if (h->key < key)
h->right = put(h->right, key, value);
else
h->value = value;
return balance(h);
}
Node* rotationRR(Node* h) {
Node* x = h->right;
h->right = x->left;
x->left = h;
return x;
}
Node* rotationLL(Node* h) {
Node* x = h->left;
h->left = x->right;
x->right = h;
return x;
}
Node* rotationLR(Node* h) {
Node* x = h->left;
h->left = rotationRR(x);
return rotationLL(h);
}
Node* rotationRL(Node* h) {
Node* x = h->right;
h->right = rotationLL(x);
return rotationRR(h);
}
Node* balance(Node* x) {
int bal_factor = diff(x);
if (bal_factor > 1) {
if (diff(x->left) > 0)
x = rotationLL(x);
else
x = rotationLR(x);
} else if (bal_factor < -1) {
if (diff(x->right) > 0)
x = rotationRL(x);
else
x = rotationRR(x);
}
return x;
}
Node* get(Node* h, K const& key) {
if (!h) return nullptr;
if (h->key < key)
return get(h->left, key);
else if (key < h->key)
return get(h->left, key);
return h;
}
int height(Node* h) {
if (!h) return 0;
return max(height(h->left), height(h->right)) + 1;
}
int diff(Node* h) { return height(h->left) - height(h->right); }
void preorder(Node* node, void (*visit)(V)) {
if (!node) return;
visit(node->key);
preorder(node->left, visit);
preorder(node->right, visit);
}
public:
AVLTree() : root(nullptr) {}
void put(K const& key, V const& value) { root = put(root, key, value); }
void preorder(void (*visit)(V)) { preorder(root, visit); }
V& get(K const& key) {
Node* x = get(root, key);
if (x == nullptr) {
put(key, V{});
x = get(root, key);
}
return x->value;
}
V& operator[](K const& key) { return get(key); }
};
|
