АВЛ-дерево

Материал из Олимпиадное программирование в УлГТУ
Перейти к навигации Перейти к поиску
class AVLTree {

    struct Node {
        int key, height;
        Node *l, *r;
        Node(int key) : key(key), height(1), l(0), r(0) {}
    } *root;

    int height(const Node *n) const {
        return n ? n->height : 0;
    }

    void updateHeight(Node *n) {
        if (n)
            n->height = 1 + max(height(n->l), height(n->r));
    }

    void simpleRotateLeft(Node *&n) {
        Node *a = n, *b = a->r;
        a->r = b->l;
        b->l = a;
        updateHeight(a);
        updateHeight(b);
        n = b;
    }

    void simpleRotateRight(Node *&n) {
        Node *a = n, *b = a->l;
        a->l = b->r;
        b->r = a;
        updateHeight(a);
        updateHeight(b);
        n = b;
    }

    void doubleRotateLeft(Node *&n) {
        simpleRotateRight(n->r);
        simpleRotateLeft(n);
    }

    void doubleRotateRight(Node *&n) {
        simpleRotateLeft(n->l);
        simpleRotateRight(n);
    }

    void rebalance(Node *&n) {
        if (!n)
            return;
        if (height(n->l) - height(n->r) == -2) {
            if (n->r && height(n->r->l) - height(n->r->r) == +1)
                doubleRotateLeft(n);
            else
                simpleRotateLeft(n);
        } 
        if (height(n->l) - height(n->r) == +2) {
            if (n->l && height(n->l->l) - height(n->l->r) == -1)
                doubleRotateRight(n);
            else
                simpleRotateRight(n);
        }
        updateHeight(n);
    }

    void insert(Node *&n, int key) {
        if (!n)
            n = new Node(key);
        else
            insert(key < n->key ? n->l : n->r, key);
        rebalance(n);
    }

    void erase(Node *&n, int key) {
        if (!n)
            return;
        if (n->key == key) {
            if (!n->l || !n->r) {
                Node *child = (n->l ? n->l : n->r);
                delete n;
                n = child;
            } else {
                Node *successor = n->r;
                while (successor->l)
                    successor = successor->l;
                n->key = successor->key;
                erase(n->r, successor->key);
            }
        } else
            erase(key < n->key ? n->l : n->r, key);
        rebalance(n);
    }

public:

    AVLTree() : root(0) {}

    int height() const {
        return height(root);
    }
   
    void insert(int key) {
        insert(root, key);
    }

    void erase(int key) {
        erase(root, key);
    }

};