C++ 进阶：AVL 树的原理与实现

存储结构

AVL 树在二叉链基础上，额外增加了父节点指针和平衡因子，属于三叉链存储结构。

#pragma once
#include <iostream>
#include <assert.h>
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 {
private:
    typedef AVLTreeNode<K, V> Node;
    Node* _root = nullptr;

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

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

    bool _IsAVLTree(Node* root) {
        if (root == nullptr) return true;
        int leftHeight = _Height(root->_left);
        int rightHeight = _Height(root->_right);
        int bf = rightHeight - leftHeight;
        if (abs(bf) >= 2 || root->_bf != bf) return false;
        return _IsAVLTree(root->_left) && _IsAVLTree(root->_right);
    }

public:
    Node* Find(const K& key) {
        Node* curr = _root;
        while (curr) {
            if (curr->_kv.first < key) curr = curr->_right;
            else if (curr->_kv.first > key) curr = curr->_left;
            else return curr;
        }
        return nullptr;
    }

    bool Insert(const pair<K, V>& kv) {
        if (_root == nullptr) {
            _root = new Node(kv);
            return true;
        }

        Node* current = _root;
        Node* parent = nullptr;

        while (current) {
            if (current->_kv.first < kv.first) {
                parent = current;
                current = current->_right;
            } else if (current->_kv.first > kv.first) {
                parent = current;
                current = current->_left;
            } else {
                return false; // Key exists
            }
        }

        current = new Node(kv);
        if (kv.first < parent->_kv.first) parent->_left = current;
        else parent->_right = current;

        current->_parent = parent;

        while (parent) {
            if (current == parent->_left) parent->_bf--;
            else parent->_bf++;

            if (parent->_bf == 0) break;
            else if (parent->_bf == -1 || parent->_bf == 1) {
                current = parent;
                parent = parent->_parent;
            } else {
                if (parent->_bf == -2 && current->_bf == -1) RotateR(parent);
                else if (parent->_bf == 2 && current->_bf == 1) RotateL(parent);
                else if (parent->_bf == -2 && current->_bf == 1) RotateLR(parent);
                else if (parent->_bf == 2 && current->_bf == -1) RotateRL(parent);
                else assert(false);
                break;
            }
        }
        return true;
    }

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

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

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

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

        subL->_bf = 0;
        parent->_bf = 0;
    }

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

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

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

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

        subR->_bf = 0;
        parent->_bf = 0;
    }

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

        RotateL(parent->_left);
        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 if (bf == 0) { subLR->_bf = 0; subL->_bf = 0; parent->_bf = 0; }
        else assert(false);
    }

    void RotateRL(Node* parent) {
        Node* subR = parent->_right;
        Node* subRL = parent->_right->_left;
        int bf = subRL->_bf;

        RotateR(parent->_right);
        RotateL(parent);

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

    void InOrder() {
        _InOrder(_root);
        cout << endl;
    }

    int Height() { return _Height(_root); }
    int Size() {
        if (_root == nullptr) return 0;
        return _Size(_root);
    }
    int _Size(Node* root) {
        if (root == nullptr) return 0;
        return _Size(root->_left) + _Size(root->_right) + 1;
    }
    bool IsAVLTree() { return _IsAVLTree(_root); }
};

测试示例

#define _CRT_SECURE_NO_WARNINGS 1
#include "AVLTree.h"

void TestAVLTree() {
    AVLTree<int, int> t;
    int a[] = { 16, 3, 7, 11, 9, 26, 18, 14, 15 };
    for (auto e : a) t.Insert({ e, e });
    t.InOrder();
    cout << "Is AVL: " << t.IsAVLTree() << endl;
}

int main() {
    TestAVLTree();
    return 0;
}