AVLTree(二叉平衡树)C++实现

链接

#ifndef ALGO_AVL_H_
#define ALGO_AVL_H_
#include 
#include  
#include 
#include 
#include 
namespace alg {

template
class AVL {
public:
    AVL():tree(0), nodeNum(0) {}
    T root() const {
        return tree->value;
    }
    unsigned height() const {
        return Node::getHeight(tree);
    }
    unsigned size() const {
        return nodeNum;
    }
    bool isEmpty() const {
        return nodeNum == 0;
    }
    bool contains(const T& x) const {
        if (!isEmpty()) {
            return tree->contain(x);
        } else 
            return false;
    }
    void insert(const T& x) {
        if (isEmpty()) 
            tree = new Node(x);
        else
            tree = tree->insert(x);
        nodeNum++;
    }
    void erase(const T& x) {
        if (!isEmpty()) {
            bool found = false;
            tree = tree->erase(x, found);
            if (found)
                nodeNum--;
        }
    }
private:
    struct Node {
        Node *left, *right;
        unsigned height;
        T value;

        Node(const T& e):value(e),left(0),right(0),height(1) {}
        bool contain(const T& e) const {
            if (e == value) 
                return true;
            else if (e < value && left != 0) 
                return left->contain(e);
            else if (e > value && right != 0) 
                return right->contain(e);
            else
                return false;
        }
        Node* insert(const T& e) {
            if (e <= value) {
                if (left == 0)
                    left = new Node(e);
                else 
                    left = left->insert(e);
            } else {
                if (right == 0)
                    right = new Node(e);
                else 
                    right = right->insert(e);
            }
            return update();
        }
        Node* erase(const T& e, bool& found) {
            if (e == value) {
                found = true;
                if (left == 0) {
                    Node* retNode = right;
                    delete this;
                    return retNode;
                } else if (right == 0) {
                    Node* retNode = left;
                    delete this;
                    return retNode;
                } else {
                    // 左子树不为空 右子树也不为空 
                    // 此时选择左子树中的最右节点(即值最大的节点)
                    // 也可以是右子树中的最左节点
                    std::stack trace;
                    Node* current = left;
                    while (current != 0) {
                        trace.push(current);
                        current = current->right;
                    }
                    // 毕竟最右端结点也有左子树
                    current = trace.top();
                    value = current->value;
                    Node* lsubtree = current->left;
                    delete current;
                    trace.pop();

                    if (trace.empty()) {
                        // 左子树没有右节点的情况
                        left = lsubtree;
                    } else {
                        trace.top()->right = lsubtree;
                        trace.pop();
                        // 更新
                        while (!trace.empty()) {
                            current = trace.top();
                            current->right = current->right->update();
                            trace.pop();
                        }
                        }
                    }
                    return update();
                } else if (e < value) {
                    if (left != 0) {
                        left = left->erase(e, found);
                        return update();
                    }
                } else {
                    if (right != 0) {
                        right = right->erase(e, found);
                        return update();
                    }
                }
                return this;
        }
        Node* update() {
            updateHeight();
            if (getBF(this) >= 2) {
                if (getBF(left) <= -1) 
                    LR();
                return LL();
            } else if(getBF(this) <= -2) {
                if (getBF(right) >= 1) 
                    RL();
                return RR();
            } else 
                return this;
        }
        void updateHeight() {
            height = std::max(getHeight(left), getHeight(right)) + 1;
        }
        void LR() {
            // 在左子树的右节点插入
            Node* originL = left;
            left = originL->right;
            originL->right = left->left;
            left->left = originL;
            originL->updateHeight();
            left->updateHeight();
            updateHeight();
        }
        void RL() {
            // 在右子树的左节点插入
            Node * originR = right;
            right = originR->left;
            originR->left = right->right;
            right->right = originR;
            originR->updateHeight();
            right->updateHeight();
            updateHeight();
        }
        Node* LL() {
            // 在左子树的左节点插入
            Node* originL = left;
            left = originL->right;
            originL->right = this;
            updateHeight();
            originL->updateHeight();
            return originL; 
        }
        Node* RR() {
            // 在左子树的左节点插入
            Node* originR = right;
            right = originR->left;
            originR->left = this;
            updateHeight();
            originR->updateHeight();
            return originR;
        }
        static int getBF(const Node* t) {
            return getHeight(t->left) - getHeight(t->right);
        }
        static int getHeight(const Node* t) {
            return t == 0 ? 0 : t->height;
        }
    };        
    Node* tree;
    unsigned nodeNum;
};
}
#endif