C++ 实现 AVL 树：原理、旋转与代码实战

探讨 AVL 树的旋转问题

当平衡因子超出 [-1, 1] 范围时，必须通过旋转来恢复平衡。旋转的核心在于保持二叉搜索树的性质（左小右大）的同时调整高度。

右单旋转

假设 10 为根的树，左侧 a 子树插入新节点导致高度增加，使得 10 的平衡因子变为 -2。此时左侧太高，需要向右旋转。

核心步骤：将 b 变成 10 的左子树，10 变成 5 的右子树，5 成为新的根。这样既符合搜索树规则，又恢复了平衡，且整棵树高度恢复到插入前的水平。

具体场景包括插入前子树高度 h=0, 1, 2, 3 等情况，抽象图可以涵盖所有具象形态。

左单旋转

逻辑与右单旋相反。若右侧太高（平衡因子为 2），则以 15 为根进行左旋，将 10 变为 15 的左子树，15 成为新根。

左右双旋

当左侧高，但新节点插入在左子树的右侧（b 子树）时，单纯右旋无法解决问题。需要先对左子树进行左旋，使其变为纯粹的左高，再对根节点进行右旋。

顺序不能调换，必须是先左后右。这涉及对 b 子树内部结构的进一步分析，根据 b 的平衡因子不同，旋转后的平衡因子赋值也有所区别。

右左双旋

同理，当右侧高且新节点插入在右子树的左侧时，需要先右旋再左旋。

AVL 树的实现层

结构定义

AVL 树节点需要包含 key-value 对、左右孩子指针、父指针以及平衡因子。

#pragma once
#include<iostream>
#include<vector>
using namespace std;
#include<assert.h>

// AVL 树的结构
template<class K,class V>
struct AVLTreeNode {
    pair<K, V> _kv;
    AVLTreeNode<K, V>* _left;
    AVLTreeNode<K, V>* _right;
    AVLTreeNode<K, V>* _parent;
    int _bf; // balance factor：平衡因子

    // 节点的构造
    AVLTreeNode(const pair<K,V>& kv) :_kv(kv), _left(nullptr), _right(nullptr), _parent(nullptr), _bf(0) {}
};

template<class K, class V>
struct AVLTree {
    typedef AVLTreeNode<K, V> Node;
public:
    // 插入函数
    bool Insert(const pair<K, V>& kv) {
        if (_root == nullptr) {
            _root = new Node(kv);
            return true;
        }
        Node* parent = nullptr;
        Node* cur = _root;
        while (cur) {
            if (cur->_kv.first < kv.first) {
                parent = cur;
                cur = cur->_right;
            } else if (cur->_kv.first > kv.first) {
                parent = cur;
                cur = cur->_left;
            } else {
                return false; // 已存在
            }
        }
        cur = new Node(kv);
        cur->_bf = 0;
        if (parent->_kv.first < kv.first) {
            parent->_right = cur;
        } else {
            parent->_left = cur;
        }
        cur->_parent = parent;

        // 控制平衡：更新平衡因子
        while (parent) {
            if (cur == parent->_left) {
                parent->_bf--;
            } else {
                parent->_bf++;
            }

            if (parent->_bf == 0) {
                break; // 高度不变，停止
            } else if (parent->_bf == 1 || parent->_bf == -1) {
                cur = parent;
                parent = parent->_parent; // 高度变了，继续向上
            } else if (parent->_bf == 2 || parent->_bf == -2) {
                // 失衡，需要旋转
                if (parent->_bf == -2 && cur->_bf == -1) {
                    RotateR(parent);
                } else if (parent->_bf == 2 && cur->_bf == 1) {
                    RotateL(parent);
                } else if (parent->_bf == -2 && cur->_bf == 1) {
                    RotateLR(parent);
                } else if (parent->_bf == 2 && cur->_bf == -1) {
                    RotateRL(parent);
                } else {
                    assert(false);
                }
                break;
            } else {
                assert(false);
            }
        }
        return true;
    }

    // 右单旋
    void RotateR(Node* parent) {
        Node* subL = parent->_left;
        Node* subLR = subL->_right;
        parent->_left = subLR;
        if (subLR) subLR->_parent = parent;
        
        Node* parentParent = parent->_parent;
        subL->_right = parent;
        parent->_parent = subL;

        if (parent == _root) {
            _root = subL;
            subL->_parent = nullptr;
        } else {
            if (parentParent->_left == parent) {
                parentParent->_left = subL;
            } else {
                parentParent->_right = subL;
            }
            subL->_parent = parentParent;
        }
        parent->_bf = subL->_bf = 0;
    }

    // 左单旋
    void RotateL(Node* parent) {
        Node* subR = parent->_right;
        Node* subRL = subR->_left;
        parent->_right = subRL;
        if (subRL) subRL->_parent = parent;

        Node* parentParent = parent->_parent;
        subR->_left = parent;
        parent->_parent = subR;

        if (parent == _root) {
            _root = subR;
            subR->_parent = nullptr;
        } else {
            if (parent == parentParent->_left) {
                parentParent->_left = subR;
            } else {
                parentParent->_right = subR;
            }
            subR->_parent = parentParent;
        }
        parent->_bf = subR->_bf = 0;
    }

    // 左右双旋
    void RotateLR(Node* parent) {
        Node* subL = parent->_left;
        Node* subLR = subL->_right;
        int bf = subLR->_bf;
        RotateL(parent->_left);
        RotateR(parent);

        if (bf == 0) {
            parent->_bf = 0; subL->_bf = 0; subLR->_bf = 0;
        } else if (bf == 1) {
            parent->_bf = 0; subL->_bf = -1; subLR->_bf = 0;
        } else if (bf == -1) {
            parent->_bf = 1; subL->_bf = 0; subLR->_bf = 0;
        } else {
            assert(false);
        }
    }

    // 右左双旋
    void RotateRL(Node* parent) {
        Node* subR = parent->_right;
        Node* subRL = subR->_left;
        int bf = subRL->_bf;
        RotateR(parent->_right);
        RotateL(parent);

        if (bf == 0) {
            subR->_bf = 0; subRL->_bf = 0; parent->_bf = 0;
        } else if (bf == 1) {
            subR->_bf = 0; subRL->_bf = 0; parent->_bf = -1;
        } else if (bf == -1) {
            subR->_bf = 1; subRL->_bf = 0; parent->_bf = 0;
        } else {
            assert(false);
        }
    }

    // 查找
    Node* Find(const K& key) {
        Node* cur = _root;
        while (cur) {
            if (cur->_kv.first < key) {
                cur = cur->_right;
            } else if (cur->_kv.first > key) {
                cur = cur->_left;
            } else {
                return cur;
            }
        }
        return nullptr;
    }

    int Size() { return _Size(_root); }
    int Height() { return _Height(_root); }
    bool IsBalanceTree() { return _IsBalanceTree(_root); }
    void InOrder() { _InOrder(_root); cout << endl; }

private:
    int _Size(Node* root) {
        if (root == nullptr) return 0;
        return _Size(root->_left) + _Size(root->_right) + 1;
    }

    int _Height(Node* root) {
        if (root == nullptr) return 0;
        int lh = _Height(root->_left);
        int rh = _Height(root->_right);
        return lh > rh ? lh + 1 : rh + 1;
    }

    bool _IsBalanceTree(Node* root) {
        if (nullptr == root) return true;
        int lh = _Height(root->_left);
        int rh = _Height(root->_right);
        if (rh - lh != root->_bf || abs(root->_bf) >= 2) {
            cout << root->_kv.first << ":平衡因子异常" << endl;
            return false;
        }
        return _IsBalanceTree(root->_left) && _IsBalanceTree(root->_right);
    }

    int _InOrder(Node* root) {
        if (root == nullptr) return 0;
        _InOrder(root->_left);
        cout << root->_kv.first << ":" << root->_kv.second << endl;
        _InOrder(root->_right);
    }

private:
    Node* _root = nullptr;
};

测试验证

通过构造特定序列测试插入和平衡性，特别是包含双旋场景的测试用例。

#define _CRT_SECURE_NO_WARNINGS 1
#include"AVLTree.h"

void TestAVLTree1() {
    AVLTree<int, int> t;
    // 特殊的带有双旋场景的测试用例
    int a[] = { 4, 2, 6, 1, 3, 5, 15, 7, 16, 14 };
    for (auto e : a) {
        t.Insert({ e,e });
        cout << e << "->" << t.IsBalanceTree() << endl;
    }
    cout << t.IsBalanceTree() << endl;
}

void TestAVLTree2() {
    const int N = 100000;
    vector<int> v;
    v.reserve(N);
    srand(time(0));
    for (size_t i = 0; i < N; i++) {
        v.push_back(rand() + i);
    }
    size_t begin2 = clock();
    AVLTree<int, int> t;
    for (auto e : v) {
        t.Insert(make_pair(e, e));
    }
    size_t end2 = clock();
    cout << "Insert:" << end2 - begin2 << endl;
    cout << t.IsBalanceTree() << endl;
    cout << "Height:" << t.Height() << endl;
    cout << "Size:" << t.Size() << endl;
    
    size_t begin1 = clock();
    for (size_t i = 0; i < N; i++) {
        t.Find((rand() + i));
    }
    size_t end1 = clock();
    cout << "Find:" << end1 - begin1 << endl;
}

int main() {
    TestAVLTree1();
    //TestAVLTree2();
    return 0;
}

运行结果会显示每次插入后的平衡状态，以及大规模数据下的性能表现。