C++手写AVL树：插入、旋转与自平衡验证

回溯更新平衡因子

插入后从新节点向根方向回溯。每步更新父节点的_bf：如果新节点是左孩子，父节点_bf加1；是右孩子，父节点_bf减1。然后根据父节点新的_bf决定下一步：

• _bf == 0：子树高度没变，不用再向上递了。
• _bf == ±1：子树高度变了，需要继续向上调整。
• _bf == ±2：失衡发生，立刻旋转，旋转后整个调整结束。

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 = cur->_parent;
    } else {
        // 失衡，准备旋转
        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;
    }
}

旋转操作

当平衡因子绝对值达到2，就根据具体形态选择旋转策略。这里用两张图先把关系捋清楚，然后再看代码会轻松很多。

右单旋（RotateR）

适用于LL型（左子树的左边过高）。基本操作就是把左孩子提上来，原根变成它的右子树，再把左孩子的右子树挂到原根的左边。

void RotateR(Node* parent) {
    Node* subL = parent->_left;
    Node* subLR = subL->_right;
    Node* ppNode = parent->_parent;

    parent->_left = subLR;
    if (subLR) subLR->_parent = parent;

    subL->_right = parent;
    parent->_parent = subL;

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

旋转后，原parent和subL都恢复了平衡，_bf置0。

左单旋（RotateL）

适用于RR型（右子树的右边过高）。逻辑和右单旋完全对称。

void RotateL(Node* parent) {
    Node* subR = parent->_right;
    Node* subRL = subR->_left;
    Node* ppNode = parent->_parent;

    parent->_right = subRL;
    if (subRL) subRL->_parent = parent;

    subR->_left = parent;
    parent->_parent = subR;

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

双旋（LR与RL）

当问题不是单边倾斜，而是'拐弯'时，比如LR型（左子树的右边高），需要先对左子树做一次左旋，变成LL型，再对当前节点做右旋。平衡因子的更新要根据中间节点（左子树的右孩子）原本的_bf值来分别处理，所以先记下它的值。

void RotateLR(Node* parent) {
    Node* subL = parent->_left;
    Node* subLR = subL->_right;
    int bf = subLR->_bf;

    RotateL(subL);
    RotateR(parent);

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

RL型对称处理即可。

其他操作

查找就是标准的BST查找，因为树高有保证，复杂度O(log N)。

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;
}

验证平衡性通常用在测试里，递归检查每个节点的平衡因子是否和高度差一致，并且绝对值不超过1。

bool _IsBalanceTree(Node* root) {
    if (!root) return true;
    int leftHeight = _Height(root->_left);
    int rightHeight = _Height(root->_right);
    int diff = leftHeight - rightHeight;

    if (abs(diff) >= 2) {
        cout << root->_kv.first << " 高度差异常" << endl;
        return false;
    }
    if (root->_bf != diff) {
        cout << root->_kv.first << " 平衡因子异常" << endl;
        return false;
    }
    return _IsBalanceTree(root->_left) && _IsBalanceTree(root->_right);
}

完整实现与测试

下面是AVLTree.h的完整代码，包含了前面所有的方法，外加两个测试函数：一个手工插入小规模数据看平衡状态，一个用百万随机数压测性能。

#pragma once
#include <iostream>
#include <vector>
#include <algorithm>
#include <cassert>
using namespace std;

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;

    AVLTreeNode(const pair<K, V>& kv)
        : _kv(kv), _left(nullptr), _right(nullptr), _parent(nullptr), _bf(0) {}
};

template <class K, class V>
class AVLTree {
    typedef AVLTreeNode<K, V> Node;
public:
    bool Insert(const pair<K, V>& kv) {
        if (!_root) {
            _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->_parent = parent;
        if (parent->_kv.first < kv.first) parent->_right = cur;
        else parent->_left = cur;

        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 = cur->_parent;
            } else {
                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;
            }
        }
        return true;
    }

    void InOrder() { _InOrder(_root); cout << endl; }
    bool IsBalanceTree() { return _IsBalanceTree(_root); }
    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 Height() { return _Height(_root); }
    int Size() { return _Size(_root); }

private:
    int _Size(Node* root) {
        if (!root) return 0;
        return _Size(root->_left) + _Size(root->_right) + 1;
    }
    int _Height(Node* root) {
        if (!root) return 0;
        int leftH = _Height(root->_left);
        int rightH = _Height(root->_right);
        return leftH > rightH ? leftH + 1 : rightH + 1;
    }
    bool _IsBalanceTree(Node* root) {
        if (!root) return true;
        int leftH = _Height(root->_left);
        int rightH = _Height(root->_right);
        int diff = leftH - rightH;
        if (abs(diff) >= 2) {
            cout << root->_kv.first << " 高度差异常" << endl;
            return false;
        }
        if (root->_bf != diff) {
            cout << root->_kv.first << " 平衡因子异常" << endl;
            return false;
        }
        return _IsBalanceTree(root->_left) && _IsBalanceTree(root->_right);
    }
    void _InOrder(Node* root) {
        if (!root) return;
        _InOrder(root->_left);
        cout << root->_kv.first << ":" << root->_kv.second << " ";
        _InOrder(root->_right);
    }
    void RotateR(Node* parent) {
        Node* subL = parent->_left;
        Node* subLR = subL->_right;
        Node* ppNode = parent->_parent;
        parent->_left = subLR;
        if (subLR) subLR->_parent = parent;
        subL->_right = parent;
        parent->_parent = subL;
        if (parent == _root) {
            subL->_parent = nullptr;
            _root = subL;
        } else {
            subL->_parent = ppNode;
            if (parent == ppNode->_right) ppNode->_right = subL;
            else ppNode->_left = subL;
        }
        parent->_bf = 0;
        subL->_bf = 0;
    }
    void RotateL(Node* parent) {
        Node* subR = parent->_right;
        Node* subRL = subR->_left;
        Node* ppNode = parent->_parent;
        parent->_right = subRL;
        if (subRL) subRL->_parent = parent;
        subR->_left = parent;
        parent->_parent = subR;
        if (parent == _root) {
            _root = subR;
            subR->_parent = nullptr;
        } else {
            subR->_parent = ppNode;
            if (parent == ppNode->_right) ppNode->_right = subR;
            else ppNode->_left = subR;
        }
        parent->_bf = 0;
        subR->_bf = 0;
    }
    void RotateLR(Node* parent) {
        Node* subL = parent->_left;
        Node* subLR = subL->_right;
        int bf = subLR->_bf;
        RotateL(subL);
        RotateR(parent);
        if (bf == 1) {
            subLR->_bf = 0; subL->_bf = 0; parent->_bf = -1;
        } else if (bf == -1) {
            subLR->_bf = 0; subL->_bf = 1; parent->_bf = 0;
        } else {
            subL->_bf = 0; subLR->_bf = 0; parent->_bf = 0;
        }
    }
    void RotateRL(Node* parent) {
        Node* subR = parent->_right;
        Node* subRL = subR->_left;
        int bf = subRL->_bf;
        RotateR(subR);
        RotateL(parent);
        if (bf == 0) {
            parent->_bf = 0; subR->_bf = 0; subRL->_bf = 0;
        } else if (bf == 1) {
            subRL->_bf = 0; parent->_bf = 0; subR->_bf = -1;
        } else if (bf == -1) {
            subRL->_bf = 0; parent->_bf = 1; subR->_bf = 0;
        }
    }
    Node* _root = nullptr;
};

// 测试用例
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 << "Insert: " << e << " -> Balance: " << t.IsBalanceTree() << endl;
    }
    t.InOrder();
    cout << "Final Balance: " << t.IsBalanceTree() << endl;
}

void TestAVLTree2() {
    const int N = 1000000;
    vector<int> v;
    v.reserve(N);
    srand((unsigned int)time(0));
    for (int i = 0; i < N; i++) v.push_back(rand() + i);

    size_t begin = clock();
    AVLTree<int, int> t;
    for (auto e : v) t.Insert(make_pair(e, e));
    size_t end = clock();

    cout << "Insert Time: " << end - begin << endl;
    cout << "Height: " << t.Height() << endl;
    cout << "Size: " << t.Size() << endl;
    cout << "Is Balanced: " << t.IsBalanceTree() << endl;
}

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

后记

AVL树用严格的平衡换来了最坏情况下仍然稳定在O(log N)的查询性能。插入时的旋转开销在查询频繁的场景里可以接受。把旋转操作的指针变动画图走一遍，搞清楚谁连到了谁，这个数据结构其实比看上去要简单。