整体代码
#define LH 1 //left high
#define EH 0 //equal high
#define RH -1 //right high
struct avl_tree {
int val;
avl_tree* left;
avl_tree* right;
int bf; //平衡因子,左边子树长度减去右边子树长度
avl_tree(int m_val, int m_bf) {
this->val = m_val;
this->bf = m_bf;
this->left = this->right = nullptr;
}
};
//右旋
avl_tree* right_rotate(avl_tree* root) {
avl_tree* left = root->left;
avl_tree* tmp = left->right;
left->right = root;
root->left = tmp;
return left;
}
//左旋
avl_tree* left_rotate(avl_tree* root) {
avl_tree* right = root->right;
avl_tree* tmp = right->left;
right->left = root;
root->right = tmp;
return right;
}
//向 root 的子树插入节点,破坏平衡时的操作
//此时 root 的 bf 大于 1
//判断 root 的 left 的 bf 是 1 还是 -1
avl_tree* left_balance(avl_tree* root) {
avl_tree* left = root->left;
switch (left->bf) {
case LH:
left->bf = root->bf = EH;
root = right_rotate(root);
break;
//当root 的 bf 和 left 的 bf 不同号时,就要双旋【先左后右】
case RH:
avl_tree* right = left->right;
switch (right->bf) {
case LH:
root->bf = RH;
left->bf = EH;
break;
case RH:
root->bf = EH;
left->bf = LH;
break;
case EH:
left->bf = root->bf = EH;
break;
}
right->bf = EH;
root->left = left_rotate(left);
root = right_rotate(root);
}
return root;
}
//右平衡与左平衡的思路类似
avl_tree* right_balance(avl_tree* root) {
avl_tree* right = root->right;
switch (right->bf) {
case RH:
root->bf = right->bf = EH;
root = left_rotate(root);
break;
case LH:
avl_tree* left = right->left;
switch (left->bf) {
case LH:
root->bf = EH;
right->bf = RH;
break;
case RH:
root->bf = LH;
right->bf = EH;
break;
//正好添加的节点的 bf 是 0
case EH:
root->bf = right->bf = 0;
}
left->bf = EH;
root->right = right_rotate(right);
root = left_rotate(root);
}
return root;
}
avl_tree* insert(avl_tree* root, int val, bool& taller) {
if (!root) {
root = new avl_tree(val, EH);
taller = true;
return root;
}
if (root->val == val) {
taller = false;
return nullptr;
}
else if (root->val < val) {
avl_tree* tmp = insert(root->right, val, taller);
if (tmp == nullptr)
return nullptr;
root->right = tmp;
//每个父节点都需要判断,为了寻找那个最近的不符合条件的结点
//找到了就将 taller 变为 false,之后的父节点就不会被处理
if (taller == true) {
switch (root->bf) {
//如果 root 左边大,因为加在右边,所以此时平衡,且以后都不会旋转,taller 改为 false
case LH:
taller = false;
root->bf = EH;
break;
//如果 root 右边大, 因为加在右边,所以此时 root 不平衡,调整左边
case RH:
root = right_balance(root);
taller = false;
break;
case EH:
root->bf = RH;
break;
}
}
}
else {
avl_tree* tmp = insert(root->left, val, taller);
if (tmp == nullptr)
return nullptr;
root->left = tmp;
if (taller == true) {
switch (root->bf) {
case LH:
root = left_balance(root);
taller = false;
break;
case RH:
taller = false;
root->bf = EH;
break;
case EH:
root->bf = LH;
break;
}
}
}
return root;
}
//前序遍历
void show(avl_tree* root) {
if (!root)
return;
cout << root->val << " ";
show(root->left);
show(root->right);
}
//左子树平衡的测试
void test_left_balance() {
avl_tree* root = new avl_tree(5, 2);
avl_tree* left1 = new avl_tree(2, -1);
avl_tree* right1 = new avl_tree(6, 0);
avl_tree* left2 = new avl_tree(1, 0);
avl_tree* right2 = new avl_tree(4, 1);
avl_tree* left3 = new avl_tree(3, 0);
root->left = left1;
root->right = right1;
left1->left = left2;
left1->right = right2;
right2->left = left3;
root = left_balance(root);
show(root);
}
//测试代码
int main() {
int i;
avl_tree* root = nullptr;
bool taller = false;;
while (cin >> i) {
root = insert(root, i, taller);
}
show(root);
return 0;
}
左旋
//左旋
avl_tree* left_rotate(avl_tree* root) {
avl_tree* right = root->right;
avl_tree* tmp = right->left;
right->left = root;
root->right = tmp;
return right;
}
/*
发生的情况是,root 的 bf 为 -2,root 的 right 的 bf 为 -1
*/
右旋
//右旋
avl_tree* right_rotate(avl_tree* root) {
avl_tree* left = root->left;
avl_tree* tmp = left->right;
left->right = root;
root->left = tmp;
return left;
}
/*
发生的情况是,root 的 bf 为 2,root 的 left 的 bf 为 1
*/
左子树平衡
/*
左子树平衡可能要右旋,可能先要左旋再右旋
左右旋发生的情况是,root 的 bf 为 2,root 的 left 的 bf 为 -1【不同号】
*/
avl_tree* left_balance(avl_tree* root) {
avl_tree* left = root->left;
switch (left->bf) {
//右旋,同号
case LH:
left->bf = root->bf = EH;
root = right_rotate(root);
break;
//左旋 + 右旋,不同号
case RH:
//获得右边的节点
avl_tree* right = left->right;
switch (right->bf) {
//在 right 的左边添加
case LH:
root->bf = RH;
left->bf = EH;
break;
//在 left 的右边添加
case RH:
root->bf = EH;
left->bf = LH;
break;
//这种情况是,right 是刚刚添加的节点【最简单的情况】
case EH:
left->bf = root->bf = EH;
break;
}
right->bf = EH;
//左旋加右旋
root->left = left_rotate(left);
root = right_rotate(root);
}
return root;
}
/*
三种情况对应的图如下
*/
右子树平衡
//右平衡与左平衡的思路类似
avl_tree* right_balance(avl_tree* root) {
avl_tree* right = root->right;
switch (right->bf) {
case RH:
root->bf = right->bf = EH;
root = left_rotate(root);
break;
case LH:
avl_tree* left = right->left;
switch (left->bf) {
case LH:
root->bf = EH;
right->bf = RH;
break;
case RH:
root->bf = LH;
right->bf = EH;
break;
case EH:
root->bf = right->bf = 0;
}
left->bf = EH;
root->right = right_rotate(right);
root = left_rotate(root);
}
return root;
}
插入函数
/*
思路:找到插入的节点位置,判断该节点在其父节点的位置
假设插入到 root 的左边
如果 root 的右边有节点,则 root 的 bf 改变,且现在已经是平衡的【高度不变】
如果 root 的左边没有节点,则 root 的 bf 变为 LH,递归判断父节点是不是平衡
root 是每层的父节点
val 是参数
taller 表示高度是不是发生变化【只会调整一次平衡】
*/
avl_tree* insert(avl_tree* root, int val, bool& taller) {
if (!root) {
root = new avl_tree(val, EH);
taller = true;
return root;
}
if (root->val == val) {
taller = false;
return nullptr;
}
else if (root->val < val) {
avl_tree* tmp = insert(root->right, val, taller);
if (tmp == nullptr)
return nullptr;
root->right = tmp;
//每个父节点都需要判断,为了寻找那个最近的不符合条件的结点
//找到了就将 taller 变为 false,之后的父节点就不会被处理
if (taller == true) {
switch (root->bf) {
//如果 root 左边大,因为加在右边,所以此时平衡,且以后都不会旋转,taller 改为 false
case LH:
taller = false;
root->bf = EH;
break;
//如果 root 右边大, 因为加在右边,所以此时 root 不平衡,调整左边
case RH:
root = right_balance(root);
taller = false;
break;
case EH:
root->bf = RH;
break;
}
}
}
else {
//思路与上面类似
avl_tree* tmp = insert(root->left, val, taller);
if (tmp == nullptr)
return nullptr;
root->left = tmp;
if (taller == true) {
switch (root->bf) {
case LH:
root = left_balance(root);
taller = false;
break;
case RH:
taller = false;
root->bf = EH;
break;
case EH:
root->bf = LH;
break;
}
}
}
return root;
}