AVLTree.hpp

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

template<class K,class V>
struct AVLTreeNode
{
    AVLTreeNode<K, V>* _left;
    AVLTreeNode<K, V>* _right; 
    AVLTreeNode<K, V>* _Parent;
    K _key;
    V _value;
    int _bf;                  //平衡因子

    AVLTreeNode(const K& key, const V& value)
        :_left(NULL)
        , _right(NULL)
        , _Parent(NULL)
        , _key(key)
        , _value(value)
        , _bf(0)
    {}
};

template<class K,class V>
class AVLTree
{
    typedef AVLTreeNode<K, V> Node;
    typedef Node* pNode;
public:
    AVLTree()
        :_pRoot(NULL)
    {}

    bool Insert(const K& key, const V& value)
    {
        //树为空
        if (NULL == _pRoot)
        {
            _pRoot = new Node(key, value);
            return true;
        }

        //树不为空，找插入位置
        pNode pCur = _pRoot;
        pNode parent = NULL;
        while (pCur)
        {
            if (key < pCur->_key)
            {
                parent = pCur;
                pCur = pCur->_left;
            }
            else if (key > pCur->_key)
            {
                parent = pCur;
                pCur = pCur->_right;
            }
            else
                return false;
        }

        //插入结点
        pCur = new Node(key, value);
        if (key < parent->_key)
            parent->_left = pCur;
        else
            parent->_right = pCur;
        pCur->_Parent = parent;

        while (parent)
        {
            //更新平衡因子
            if (parent->_left == pCur)
                parent->_bf--;
            else
                parent->_bf++;

            //平衡因子有三种情况 
            //1、平衡因子为0，直接返回，不做调整 
            if (0 == parent->_bf)
                return true;
            //2、平衡因子为1或-1，向上调整平衡因子，直到平衡 
            else if (-1 == parent->_bf || 1 == parent->_bf)
            {
                pCur = parent;
                parent = pCur->_Parent;
            }
            //3、平衡因子为2或-2，进行旋转处理
            else
            {
                if (2 == parent->_bf)
                {
                    if (1 == pCur->_bf)
                    {
                        _RotateL(parent);
                    }
                    else
                    {
                        _RotateRL(parent);
                    }
                }
                else
                {
                    if (-1 == pCur->_bf)
                    {
                        _RotateR(parent);
                    }
                    else
                    {
                        _RotateLR(parent);
                    }
                }
                break;
            }
        }
        return true;
    }

    void Delete(const K& key, const V& value)
    {
        //树为空
        if (NULL == _pRoot)
            return false;

        //找要删除的结点
        pNode pDel = _pRoot;   //待删除的结点
        pNode parent = NULL;
        while (pDel)
        {
            if (key < pDel->_key)
            {
                parent = pDel;
                pCur = pDel->_left;
            }
            else if (key>pDel->_key)
            {
                parent = pDel;
                pCur = pDel->_right;
            }
            else if (key == pDel->_key)
                  break;
            else
                return false;    
        }
        //删除结点

        if (parent->_left == pDel)
        {
            delete pDel;
            parent->_left = NULL;
            parent->_bf++;
        }
        else
        {
            delete pDel;
            parent->_right = NULL;
            parent->_bf--;
        }

        //调节平衡因子
        pNode pCur = parent;
        parent = parent->_Parent;
        while (parent)
        {
            if (parent->_left == pCur)
                parent->_bf++;
            else
                parent->_bf--;

            //平衡因子有三种情况 
            //1、平衡因子为0，向上调整平衡因子，直到平衡 
            if (0 == parent->_bf)
            {
                pCur = parent;
                parent = pCur->_Parent;
            }
            //2、平衡因子为1或-1，直接返回 
            else if (-1 == parent->_bf || 1 == parent->_bf)
                    return true;
            //3、平衡因子为2或-2，进行旋转处理
            else
            {
                if (2 == parent->_bf)
                {
                    if (1 == pCur->_bf)
                    {
                        _RotateL(parent);
                    }
                    else
                    {
                        _RotateRL(parent);
                    }
                }
                else
                {
                    if (-1 == pCur->_bf)
                    {
                        _RotateR(parent);
                    }
                    else
                    {
                        _RotateLR(parent);
                    }
                }
            }
        }
    }


    void PreOrder()
    {
        cout << "先序遍历结果：";
        _PreOrder(_pRoot);
        cout << endl;
    }

    void InOrder()
    {
        cout << "中序遍历结果：";
        _InOrder(_pRoot);
        cout << endl;
    }

    size_t Height()
    {
        return _Height(_pRoot);
    }

    bool IsBalanceTree()
    {
        return _IsBalanceTree(_pRoot);
    }

private:

    bool _IsBalanceTree(pNode pRoot)
    {
        if (NULL == pRoot)
            return true;

        size_t HL = _Height(pRoot->_left);
        size_t HR = _Height(pRoot->_right);

        int tmp = HL - HR;
        if (abs(tmp) == pRoot->_bf)
            return true;
        return false;
    }


    size_t _Height(pNode pRoot)
    {
        if (pRoot == NULL)
            return 0;

        size_t hl = _Height(pRoot->_left);
        size_t hr = _Height(pRoot->_right);

        return hl > hr ? (hl + 1) : (hr + 1);
    }

    void _PreOrder(pNode pRoot)
    {
        if (pRoot)
        {
            cout << pRoot->_key << " ";
            _PreOrder(pRoot->_left);
            _PreOrder(pRoot->_right);
        }
    }

    void _InOrder(pNode pRoot)
    {
        if (pRoot)
        {
            _InOrder(pRoot->_left);
            cout << pRoot->_key << " ";
            _InOrder(pRoot->_right);
        }
    }

    void _RotateL(pNode parent)
    {
        pNode pSubR = parent->_right;
        pNode pSubRL = pSubR->_left;

        parent->_right = pSubRL;

        if (pSubRL)
            pSubRL->_Parent = parent;

        pSubR->_left = parent;
        pNode pParent = parent->_Parent;
        parent->_Parent = pSubR;
        pSubR->_Parent = pParent;

        if (pParent == NULL)
        {
            _pRoot = pSubR;
            pSubR->_Parent = NULL;
        }
        else
        {
            if (pParent->_left == parent)
                pParent->_left = pSubR;
            else
                pParent->_right = pSubR;
        }

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

    void _RotateR(pNode parent)
    {
        pNode pSubL = parent->_left;
        pNode pSubLR = pSubL->_right;
        parent->_left = pSubLR;
        if (pSubLR)
            pSubLR->_Parent = parent;

        pSubL->_right = parent;
        pNode pParent = parent->_Parent;
        parent->_Parent = pSubL;
        pSubL->_Parent = pParent;

        if (pParent == NULL)
        {
            _pRoot = pSubL;
            pSubL->_Parent = NULL;
        }
        else
        {
            if (pParent->_left == parent)
                pParent->_left = pSubL;
            else
                pParent->_right = pSubL;
        }
        parent->_bf = pSubL->_bf = 0;
    }

    void _RotateRL(pNode parent)
    {
        pNode pSubR = parent->_right;
        pNode pSubRL = pSubR->_left;

        int bf = pSubRL->_bf;
        _RotateR(pSubR);
        _RotateL(parent);

        if (bf == 1)
            parent->_bf = -1;
        else if (bf == -1)
            pSubR->_bf = 1;
    }

    void _RotateLR(pNode parent)
    {
        pNode pSubL = parent->_left;
        pNode pSubLR = pSubL->_right;
        int bf = pSubLR->_bf;

        _RotateL(pSubL);
        _RotateR(parent);

        if (bf == 1)
            pSubL->_bf = -1;
        else if (bf == -1)
            parent->_bf = 1;
    }

private:
    pNode _pRoot;
};

Test.cpp

#include "AVLTree.h"
void TestAVLTree()
{
    int array[] = {16, 3, 7, 11, 9, 26, 18, 14, 15};
    //int array[] = { 10, 20, 30, 40 };
    AVLTree<int, int> t;
    for (size_t i = 0; i < sizeof(array) / sizeof(array[0]); ++i)
        t.Insert(array[i], i);

    t.PreOrder();
    t.InOrder();

    cout <<"树的高度："<< t.Height() << endl;

    if (t.IsBalanceTree())
        cout << "是 AVL 树" << endl;
    else
        cout << "不是 AVL 树" << endl;
}

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