数据结构之红黑树的实现

#include<iostream>
using namespace std;
#include<assert.h>

enum color{
    RED,
    BLACK,
};
template<class K, class V>
struct TreeNode{
public:
    TreeNode<K, V>* _left;
    TreeNode<K, V>* _right;
    TreeNode<K, V>* _parent;
    color _color;
    V _val;;
    TreeNode(const V& val)
      :_val(val)
      , _left(NULL)
      , _right(NULL)
      , _parent(NULL)
      , _color(RED)
    {}
};
template<class K, class V>
class RBTree{
public:
    typedef TreeNode<K, V> Node;
    RBTree()
        :_root(NULL)
    {}
    bool Insert(const V& val){
        if (_root == NULL){
            _root = new Node(val);
            _root->_color = BLACK;
            return true;
        }
        Node* cur = _root;
        Node* parent = NULL;
        while (cur){
            if (val < cur->_val){
                parent = cur;
                cur = cur->_left;
            }
            else if (val>cur->_val){
                parent = cur;
                cur = cur->_right;
            }
            else{
                return false;
            }
        }
        cur = new Node(val);
        if (cur->_val < parent->_val){
            parent->_left = cur;
            cur->_parent = parent;
        }
        else{
            parent->_right = cur;
            cur->_parent = parent;
        }
        //调整颜色
        Node* grandfather = parent->_parent;
        while (parent&&parent->_color == RED){
            if (grandfather->_left == parent){
                Node* uncle = grandfather->_right;
                if (uncle&&uncle->_color == RED){//叔叔存在而且是红色
                    parent->_color = BLACK;
                    uncle->_color = BLACK;
                    grandfather->_color = RED;
                    cur = grandfather;
                    parent = cur->_parent;
                }
                else{//叔叔为黑色/叔叔不存在
                    if (parent->_right == cur){
                        RotateL(parent);
                        swap(cur, parent);
                    }
                    RotateR(grandfather);
                    parent->_color = BLACK;
                    grandfather->_color= RED;
                    break;
                }
            }
            else{
                Node* uncle = grandfather->_left;
                if (uncle&&uncle->_color == RED){
                    uncle->_color = BLACK;
                    parent->_color = BLACK;
                    cur = grandfather;
                    parent = cur->_parent;
                }
                else{
                    if (parent->_left == cur){
                        RotateR(parent);
                    swap(cur, parent);
                    }
                    RotateL(grandfather);
                    parent->_color = BLACK;
                    grandfather->_color = RED;
                    break;
                }
            }
        }
        _root->_color = BLACK;
        return true;
    }
    void RotateL(Node* parent){
        assert(parent);
        Node* ppNode = parent->_parent;
        Node* subR = parent->_right;
        Node* subRL = subR->_left;
        parent->_left = subRL;
        if (subRL){
            subRL->_parent = parent;
        }
        if (ppNode == NULL){
            _root = subR;
            subR->_parent = ppNode;
        }
        else{
            if (ppNode->_left == parent){
                ppNode->_left = subR;
                subR->_parent = ppNode;
            }
            else{
                ppNode->_right = subR;
                subR->_parent = ppNode;
            }
        }
        subR->_left = parent;
        parent->_parent = subR;
    }
    void RotateR(Node* parent){
        assert(parent);
        Node* ppNode = parent->_parent;
        Node* subL = parent->_left;
        Node* subLR = subL->_right;
        parent->_right = subLR;
        if (subLR){
            subLR->_parent = parent;
        }
        if (ppNode == NULL){
            _root = subL;
            subL->_parent = ppNode;
        }
        else{
            if (ppNode->_left == parent){
                ppNode->_left = subL;
                subL->_parent = ppNode;
            }
            else{
                ppNode->_right = subL;
                subL->_parent = ppNode;
            }
        }
        subL->_right = parent;
        parent->_parent = subL;
    }
    bool IsBlance(){
        if (_root&&_root->_color == RED){
            return false;
        }
        int N = 0;int  blacknum = 0;
        return _IsBalance(_root, blacknum, N);
    }
    bool _IsBalance(Node* root, size_t blacknum, const size_t N){
        if (root == NULL){
            return true;
        }
        if (root->_color == BLACK){
            ++blacknum;
        }
        if (root->_color == RED&&root->_parent->_color == RED){
            cout << "颜色异常" << endl;
        }
        return _IsBalance(root->_left, blacknum, N) && _IsBalance(root->_right, blacknum, N);
    }
private:
    Node* _root;
};

你可能感兴趣的:(数据结构)