AVL树

数据结构上机课花了点时间实现的AVL树，基本上在抄书。。

包含了二叉树、二叉查找树和AVL树的实现，不过毕竟没有真正学过C++，对OOP也还不算很熟悉。。碰到了一些问题：

• 继承的类也不能访问父类的private成员，不过可以用using 父类::成员名的方式访问父类的protected成员
• 子类中重载了父类的某个成员函数后，对子类和父类中的该函数同时加virtual关键字修饰，可以在运行时判断具体需要调用的函数是哪一个版本

还没有深入研究过虚函数……总之OOP有点复杂，但好像有点优美？

#include <iostream>
#include <stack>
using namespace std;

enum Balance_factor {left_higher, equal_height, right_higher};
enum Error_code {success, not_present, duplicate_error};

template <class Record>
struct Binary_node
{
    Record data;
    Binary_node<Record> *left, *right;

    Binary_node() {left = right = NULL;};
    Binary_node(const Record &x) {data = x; left = right = NULL;};

    virtual void set_balance(Balance_factor b) {};
    virtual Balance_factor get_balance() const {return equal_height;};
};

template <class Record>
class Binary_tree
{
public:
    Binary_tree() {root = NULL; count = 0;};
    bool empty() const {return !root;};
    int size() const {return count;};
    int height() const
    {
        if (!count) return 0;
        int tmp, i;
        for (tmp = 1, i = 0; tmp <= count; ++i) tmp <<= 1;
        return i;
    }
    void preorder(void (*visit)(Record &)) {recursive_preorder(root, visit);};
    void inorder(void (*visit)(Record &)) {recursive_inorder(root, visit);};
    void postorder(void (*visit)(Record &)) {recursive_postorder(root, visit);};
    void insert(Record &);
protected:
    Binary_node<Record> *root;
    int count;
    void recursive_preorder(Binary_node<Record> *sub_root, void (*visit)(Record &))
    {
        if (sub_root)
        {
            (*visit)(sub_root->data);
            recursive_preorder(sub_root->left, visit);
            recursive_preorder(sub_root->right, visit);
        }
    }
    void recursive_inorder(Binary_node<Record> *sub_root, void (*visit)(Record &))
    {
        if (sub_root)
        {
            recursive_inorder(sub_root->left, visit);
            (*visit)(sub_root->data);
            recursive_inorder(sub_root->right, visit);
        }
    }
    void recursive_postorder(Binary_node<Record> *sub_root, void (*visit)(Record &))
    {
        if (sub_root)
        {
            recursive_postorder(sub_root->left, visit);
            recursive_postorder(sub_root->right, visit);
            (*visit)(sub_root->data);
        }
    }
};

template <class Record>
void Binary_tree<Record>::insert(Record &x)
{
    if(empty())
    {
        root = new Binary_node<Record>(x);
        ++count;
        return;
    }
    stack<int> numbers;
    int item = 0, tmpcount = size();
    while (tmpcount > 0)
    {
        numbers.push((tmpcount&1) ? 1:2);
        tmpcount = (tmpcount-1)>>1;
    }
    Binary_node<Record> *current = root;
    while (numbers.size() > 1)
    {
        item = numbers.top();
        if (item == 1) current = current->left;
        if (item == 2) current = current->right;
        numbers.pop();
    }
    item = numbers.top();
    if (item == 1) current->left = new Binary_node<Record>(x);
    if (item == 2) current->right = new Binary_node<Record>(x);
    ++count;
}

template <class Record>
class Search_tree: public Binary_tree<Record>
{
public:
    Error_code insert(const Record &new_data)
    {
        Error_code result = search_and_insert(root, new_data);
        if (result == success) ++count;
        return result;
    }
    Error_code remove(const Record &target)
    {
        Error_code result = search_and_destroy(root, target);
        if (result == success) --count;
        return result;
    }
    Error_code tree_search(Record &target) const
    {
        Error_code result = success;
        Binary_node<Record> *found = search_for_node(root, target);
        if (!found) result = not_present;
        else target = found->data;
        return result;
    }
protected:
    using Binary_tree<Record>::root;
    using Binary_tree<Record>::count;

    Binary_node<Record> *search_for_node(Binary_node<Record>* sub_root, const Record &target) const;
    Error_code search_and_insert(Binary_node<Record> * &sub_root, const Record &new_data);
    Error_code search_and_destroy(Binary_node<Record>* &sub_root, const Record &target);
    Error_code remove_root(Binary_node<Record> * &sub_root);
};

template <class Record>
Binary_node<Record> *Search_tree<Record>::search_for_node(
Binary_node<Record>* sub_root, const Record &target) const
{
    if (!sub_root || sub_root->data == target)
        return sub_root;
    else if (sub_root->data < target)
        return search_for_node(sub_root->right, target);
    else return search_for_node(sub_root->left, target);
}

template <class Record>
Error_code Search_tree<Record>::search_and_insert(
Binary_node<Record> * &sub_root, const Record &new_data)
{
    if (!sub_root)
    {
        sub_root = new Binary_node<Record>(new_data);
        return success;
    }
    else if (new_data < sub_root->data)
        return search_and_insert(sub_root->left, new_data);
    else if (new_data > sub_root->data)
        return search_and_insert(sub_root->right, new_data);
    else return duplicate_error;
}

template <class Record>
Error_code Search_tree<Record>::remove_root(Binary_node<Record> * &sub_root)
{
    if (!sub_root) return not_present;
    Binary_node<Record> *to_delete = sub_root;
    if (!sub_root->right) sub_root = sub_root->left;
    else if (!sub_root->left) sub_root = sub_root->right;
    else {
        to_delete = sub_root->left;
        Binary_node<Record> *parent = sub_root;
        while (to_delete->right)
        {
            parent = to_delete;
            to_delete = to_delete->right;
        }
        sub_root->data = to_delete->data;
        if (parent == sub_root) sub_root->left = to_delete->left;
        else parent->right = to_delete->left;
    }
    delete to_delete;
    return success;
}

template <class Record>
Error_code Search_tree<Record>::search_and_destroy(
Binary_node<Record>* &sub_root, const Record &target)
{
    if (!sub_root || sub_root->data == target)
        return remove_root(sub_root);
    else if (target < sub_root->data)
        return search_and_destroy(sub_root->left, target);
    else
        return search_and_destroy(sub_root->right, target);
}

template <class Record>
struct AVL_node: public Binary_node<Record>
{
    using Binary_node<Record>::left;
    using Binary_node<Record>::right;
    using Binary_node<Record>::data;

    Balance_factor balance;

    AVL_node() {left = right = NULL; balance = equal_height;};
    AVL_node(const Record &x)
    {
        data = x;
        left = right = NULL;
        balance = equal_height;
    };

    void set_balance(Balance_factor b) {balance = b;};
    Balance_factor get_balance() const {return balance;};
};

template <class Record>
class AVL_tree: public Search_tree<Record>
{
public:
    Error_code insert(const Record &new_data)
    {
        bool taller;
        return avl_insert(root, new_data, taller);
    }
    Error_code remove(Record &new_data)
    {
        bool shorter = true;
        return avl_remove(root, new_data, shorter);
    };
private:
    using Binary_tree<Record>::root;

    Error_code avl_insert(Binary_node<Record> * &sub_root, const Record &new_data, bool &taller);
    void rotate_left(Binary_node<Record> * &sub_root);
    void rotate_right(Binary_node<Record> * &sub_root);
    void right_balance(Binary_node<Record> * &sub_root);
    void left_balance(Binary_node<Record> * &sub_root);

    Error_code avl_remove(Binary_node<Record> * &sub_root, Record &new_data, bool &shorter);
    bool right_balance2(Binary_node<Record> * &sub_root);
    bool left_balance2(Binary_node<Record> * &sub_root);
};

template <class Record>
Error_code AVL_tree<Record>::avl_insert(Binary_node<Record> * &sub_root,
const Record &new_data, bool &taller)
{
    Error_code result = success;
    if (!sub_root)
    {
        sub_root = new AVL_node<Record>(new_data);
        taller = true;
    }
    else if (new_data == sub_root->data)
    {
        result = duplicate_error;
        taller = false;
    }
    else if (new_data < sub_root->data)
    {
        result = avl_insert(sub_root->left, new_data, taller);
        if (taller)
        switch (sub_root->get_balance())
        {
        case left_higher:
            left_balance(sub_root);
            taller = false;
            break;
        case equal_height:
            sub_root->set_balance(left_higher);
            break;
        case right_higher:
            sub_root->set_balance(equal_height);
            taller = false;
            break;
        }
    }
    else
    {
        result = avl_insert(sub_root->right, new_data, taller);
        if (taller)
        switch (sub_root->get_balance())
        {
        case left_higher:
            sub_root->set_balance(equal_height);
            taller = false;
            break;
        case equal_height:
            sub_root->set_balance(right_higher);
            break;
        case right_higher:
            right_balance(sub_root);
            taller = false;
            break;
        }
    }
    return result;
}

template <class Record>
void AVL_tree<Record>::rotate_left(Binary_node<Record> * &sub_root)
{
    if (!sub_root || !sub_root->right)
        cout << "WARNING: program error detected in rotate left" << endl;
    else
    {
        Binary_node<Record> *right_tree = sub_root->right;
        sub_root->right = right_tree->left;
        right_tree->left = sub_root;
        sub_root = right_tree;
    }
}

template <class Record>
void AVL_tree<Record> :: rotate_right(Binary_node<Record> * &sub_root)
{
    if (!sub_root || !sub_root->left)
        cout << "WARNING: program error detected in rotate right" << endl;
    else
    {
        Binary_node<Record> *left_tree = sub_root->left;
        sub_root->left = left_tree->right;
        left_tree->right = sub_root;
        sub_root = left_tree;
    }
}

template <class Record>
void AVL_tree<Record>::right_balance(Binary_node<Record> * &sub_root)
{
    Binary_node<Record> * &right_tree = sub_root->right;
    switch (right_tree->get_balance())
    {
    case right_higher:
        sub_root->set_balance(equal_height);
        right_tree->set_balance(equal_height);
        rotate_left(sub_root);
        break;
    case equal_height:
        cout << "WARNING: program error in right balance" << endl;
    case left_higher:
        Binary_node<Record> *sub_tree = right_tree->left;
        switch (sub_tree->get_balance())
        {
        case equal_height:
            sub_root->set_balance(equal_height);
            right_tree->set_balance(equal_height);
            break;
        case left_higher:
            sub_root->set_balance(equal_height);
            right_tree->set_balance(right_higher);
            break;
        case right_higher:
            sub_root->set_balance(left_higher);
            right_tree->set_balance(equal_height);
            break;
        }
        sub_tree->set_balance(equal_height);
        rotate_right(right_tree);
        rotate_left(sub_root);
        break;
    }
}

template <class Record>
void AVL_tree<Record>::left_balance(Binary_node<Record> * &sub_root)
{
    Binary_node<Record> * &left_tree = sub_root->left;
    switch (left_tree->get_balance())
    {
    case left_higher:
        sub_root->set_balance(equal_height);
        left_tree->set_balance(equal_height);
        rotate_right(sub_root);
        break;
    case equal_height:
        cout << "WARNING: program error in right balance" << endl;
    case right_higher:
        Binary_node<Record> *sub_tree = left_tree->right;
        switch (sub_tree->get_balance())
        {
        case equal_height:
            sub_root->set_balance(equal_height);
            left_tree->set_balance(equal_height);
            break;
        case right_higher:
            sub_root->set_balance(equal_height);
            left_tree->set_balance(left_higher);
            break;
        case left_higher:
            sub_root->set_balance(right_higher);
            left_tree->set_balance(equal_height);
            break;
        }
        sub_tree->set_balance(equal_height);
        rotate_left(left_tree);
        rotate_right(sub_root);
        break;
    }
}

template <class Record>
Error_code AVL_tree<Record>::avl_remove(Binary_node<Record> * &sub_root,
Record &new_data, bool &shorter)
{
    Error_code result = success;
    Record sub_record;
    if (!sub_root)
    {
        shorter = false;
        return not_present;
    }
    else if (new_data == sub_root->data)
    {
        Binary_node<Record> *to_delete = sub_root;
        if (!sub_root->right)
        {
            sub_root = sub_root->left;
            shorter = true;
            delete to_delete;
            return success;
        }
        else if (!sub_root->left)
        {
            sub_root = sub_root->right;
            shorter = true;
            delete to_delete;
            return success;
        }
        else
        {
            to_delete = sub_root->left;
            Binary_node<Record> *parent = sub_root;
            while (to_delete->right)
            {
                parent = to_delete;
                to_delete = to_delete->right;
            }
            new_data = to_delete->data;
            sub_record = new_data;
        }
    }
    if (new_data < sub_root->data)
    {
        result = avl_remove(sub_root->left, new_data, shorter);
        if (sub_record.the_key()) sub_root->data = sub_record;
        if (shorter)
        switch (sub_root->get_balance())
        {
        case left_higher:
            sub_root->set_balance(equal_height);
            break;
        case equal_height:
            sub_root->set_balance(right_higher);
            shorter = false;
            break;
        case right_higher:
            shorter = right_balance2(sub_root);
            break;
        }
    }
    if (new_data > sub_root->data)
    {
        result = avl_remove(sub_root->right, new_data, shorter);
        if (sub_record.the_key()) sub_root->data = sub_record;
        if (shorter)
        switch (sub_root->get_balance())
        {
        case left_higher:
            shorter = left_balance2(sub_root);
            break;
        case equal_height:
            sub_root->set_balance(left_higher);
            shorter = false;
            break;
        case right_higher:
            sub_root->set_balance(equal_height);
            break;
        }
    }
    return result;
}

template <class Record>
bool AVL_tree<Record>::right_balance2(Binary_node<Record> * &sub_root)
{
    bool shorter;
    Binary_node<Record> * &right_tree = sub_root->right;
    switch (right_tree->get_balance())
    {
    case right_higher:
        sub_root->set_balance(equal_height);
        right_tree->set_balance(equal_height);
        rotate_left(sub_root);
        shorter = true;
        break;
    case equal_height:
        right_tree->set_balance(left_higher);
        rotate_left(sub_root);
        shorter = false;
        break;
    case left_higher:
        Binary_node<Record> *sub_tree = right_tree->left;
        switch (sub_tree->get_balance())
        {
        case equal_height:
            sub_root->set_balance(equal_height);
            right_tree->set_balance(equal_height);
            break;
        case left_higher:
            sub_root->set_balance(equal_height);
            right_tree->set_balance(right_higher);
            break;
        case right_higher:
            sub_root->set_balance(left_higher);
            right_tree->set_balance(equal_height);
            break;
        }
        sub_tree->set_balance(equal_height);
        rotate_right(right_tree);
        rotate_left(sub_root);
        shorter = true;
        break;
    }
    return shorter;
}

template <class Record>
bool AVL_tree<Record>::left_balance2(Binary_node<Record> * &sub_root)
{
    bool shorter;
    Binary_node<Record> * &left_tree = sub_root->left;
    switch (left_tree->get_balance())
    {
    case left_higher:
        sub_root->set_balance(equal_height);
        left_tree->set_balance(equal_height);
        rotate_right(sub_root);
        shorter = true;
        break;
    case equal_height:
        left_tree->set_balance(right_higher);
        rotate_right(sub_root);
        shorter = false;
        break;
    case right_higher:
        Binary_node<Record> *sub_tree = left_tree->right;
        switch (sub_tree->get_balance())
        {
        case equal_height:
            sub_root->set_balance(equal_height);
            left_tree->set_balance(equal_height);
            break;
        case right_higher:
            sub_root->set_balance(equal_height);
            left_tree->set_balance(left_higher);
            break;
        case left_higher:
            sub_root->set_balance(right_higher);
            left_tree->set_balance(equal_height);
            break;
        }
        sub_tree->set_balance(equal_height);
        rotate_left(left_tree);
        rotate_right(sub_root);
        shorter = true;
        break;
    }
    return shorter;
}

template <class Record>
void print(Record &x)
{
    cout << x << " ";
}

typedef char Record;

int main()
{
    AVL_tree<Record> mytree;
    mytree.insert('A');
    mytree.insert('V');
    mytree.insert('L');
    mytree.insert('T');
    mytree.insert('R');
    mytree.insert('E');
    mytree.insert('I');
    mytree.insert('S');
    mytree.insert('O');
    mytree.insert('K');

    cout << "Preorder:" << endl;
    mytree.preorder(print);
    cout << endl << endl;
    cout << "Inorder:" << endl;
    mytree.inorder(print);
    cout << endl << endl;
    cout << "Postorder:" << endl;
    mytree.postorder(print);
    cout << endl << endl;

    cin.get();
    return 0;
}