红黑树的实现（C++）

// 红黑树的实现，详见《算法导论第二版》第163页。

#ifndef RBT_H
#define RBT_H

namespace NTCI
{
    template<typename Type>
    class RBT
    {
    public:
        enum Color {RED, BLACK};
        class Node
        {
        public:
            Node() : color(BLACK), left(0), right(0), parent(0) {}
            Type item;
            Color color;
            Node* left;
            Node* right;
            Node* parent;
        };
        typedef Node* Link;

    public:
        RBT(Link r = Nil()) : root(r) {}
        template<typename InputIterator>
        RBT(InputIterator start, InputIterator end) : root(Nil())
        {
            for (InputIterator i = start; i != end; ++i)
                Insert(*i);
        }

        virtual ~RBT()
        {
            Clear();
        }

        static Link Nil()
        {
            static Node nil;

            return &nil;
        }

        void LeftRotate(Link link)
        {
            Link subr = link->right;
            if (subr == Nil())
                throw std::invalid_argument("待左旋的结点其右子树必须不为空");

            subr->parent = link->parent;
            if (link->parent == Nil())
                root = subr;
            else if (link == link->parent->left)
                link->parent->left = subr;
            else
                link->parent->right = subr;

            link->right = subr->left;
            if (subr->left != Nil())
                subr->left->parent = link;

            subr->left = link;
            link->parent = subr;
        }

        void RightRotate(Link link)
        {
            Link subl = link->left;
            if (subl == Nil())
                throw std::invalid_argument("待右旋的结点其左子树必须不为空");

            subl->parent = link->parent;
            if (link->parent == Nil())
                root = subl;
            else if (link == link->parent->left)
                link->parent->left = subl;
            else
                link->parent->right = subl;

            link->left = subl->right;
            if (subl->right != Nil())
                subl->right->parent = link;

            subl->right = link;
            link->parent = subl;
        }

        void Insert(Type item)
        {
            Link n = new Node();
            n->item = item;
            n->color = RED;
            n->left = Nil();
            n->right = Nil();
            n->parent = Nil();

            InsertRBT(root, n);
            RBInsertFixup(n);
        }

        void RBInsertFixup(Link link)
        {
            while (link->parent->color == RED)
            {
                if (link->parent == link->parent->parent->left)
                {
                    Link runcle = link->parent->parent->right;

                    if (runcle->color == RED) // Case1

                    {
                        link->parent->color = BLACK;
                        runcle->color = BLACK;
                        link->parent->parent->color = RED;
                        link = link->parent->parent;
                    }
                    else
                    {
                        if (link == link->parent->right) // Case2

                        {
                            link = link->parent;
                            LeftRotate(link);
                        }

                        link->parent->color = BLACK;
                        link->parent->parent->color = RED;
                        RightRotate(link->parent->parent);
                    }
                }
                else
                {
                    Link luncle = link->parent->parent->left;

                    if (luncle->color == RED) // Case1

                    {
                        link->parent->color = BLACK;
                        luncle->color = BLACK;
                        link->parent->parent->color = RED;
                        link = link->parent->parent;
                    }
                    else
                    {
                        if (link == link->parent->left) // Case2

                        {
                            link = link->parent;
                            RightRotate(link);
                        }
                        // Case3

                        link->parent->color = BLACK;
                        link->parent->parent->color = RED;
                        LeftRotate(link->parent->parent);
                    }
                }
            }

            root->color = BLACK;
        }

        Link Remove(Link link)
        {
            // 删除空节点时抛出参数无效的异常。

            if (link == Nil())
                throw std::invalid_argument("待删除的结点为空");

            // 先确定将要被删除的结点。如果结点只有一个或零个子结点，则原元素将被删除，

            // 否侧其后继结点将被删除。

            Link removed;
            if (link->left == Nil() || link->right == Nil())
                removed = link;
            else
                removed = Successor(link);

            // 待删除的元素至多只有一个子结点。

            Link child;
            if (link->left != Nil())
                child = removed->left;
            else
                child = removed->right;

            // 先设置子结点的新的父节点。

            child->parent = removed->parent;

            // 如果待删除的结点为根节点，直接将根结点设置为子结点，

            if (removed->parent == Nil())
                root = child;
            // 否则将待删除结点的父节点的子结点指向待删除结点的子结点。

            else if (removed == removed->parent->left)
                removed->parent->left = child;
            else
                removed->parent->right = child;

            // 如果原结点和被删除的结点不同，则说明原结点具有两个子结点，此时需要将

            // 被删除结点的值复制到原结点。

            if (link != removed)
                link->item = removed->item;

            if (removed->color == BLACK)
                RBRemoveFixup(child);
            return removed;
        }

        void RBRemoveFixup(Link link)
        {
            while (link != root && link->color == BLACK)
            {
                if (link == link->parent->left)
                {
                    Link rbro = link->parent->right;
                    if (rbro->color == RED) // Case1

                    {
                        rbro->color = BLACK;
                        link->parent->color = RED;
                        LeftRotate(link->parent);
                        rbro = link->parent->right;
                    }
                    if (rbro->left->color == BLACK && rbro->right->color == BLACK) // Case2

                    {
                        rbro->color = RED;
                        link = link->parent;
                    }
                    else
                    {
                        if (rbro->right->color == BLACK)
                        {
                            rbro->left->color = BLACK;
                            rbro->color = RED;
                            RightRotate(rbro);
                            rbro = link->parent->right;
                        }
                        rbro->color = link->parent->color;
                        link->parent->color = BLACK;
                        rbro->right->color = BLACK;
                        LeftRotate(link->parent);
                        link = root;
                    }
                }
                else
                {
                    Link lbro = link->parent->left;
                    if (lbro->color == RED) // Case1

                    {
                        lbro->color = BLACK;
                        link->parent->color = RED;
                        RightRotate(link->parent);
                        lbro = link->parent->left;
                    }
                    if (lbro->right->color == BLACK && lbro->left->color == BLACK) // Case2

                    {
                        lbro->color = RED;
                        link = link->parent;
                    }
                    else
                    {
                        if (lbro->left->color == BLACK)
                        {
                            lbro->right->color = BLACK;
                            lbro->color = RED;
                            LeftRotate(lbro);
                            lbro = link->parent->left;
                        }
                        lbro->color = link->parent->color;
                        link->parent->color = BLACK;
                        lbro->left->color = BLACK;
                        RightRotate(link->parent);
                        link = root;
                    }
                }
            }
            link->color = BLACK;
        }

        Link Remove(Type item)
        {
            Link link = Search(item);
            return Remove(link);
        }

        Link Search(Type item)
        {
            return SearchRBT(root, item);
        }

        int Maximum()
        {
            Link max = MaximumRBT(root);
            if (max == 0)
                throw std::out_of_range("红黑树为空");
            return max->item;
        }

        int Minimun()
        {
            Link min = MinimunRBT(root);
            if (min == 0)
                throw std::out_of_range("红黑树为空");
            return min->item;
        }

        void InorderWalk(std::ostream& out)
        {
            InorderWalkRBT(root, out);
        }

        Link Successor(Link link)
        {
            // 如果右子结点不为空，则返回右子树的最小结点，

            if (link->right != Nil())
                return MinimunRBT(link->right);

            // 否则后继元素是一个祖先结点，该结点的左子结点也是祖先结点。

            Link p = link->parent;
            while (p != Nil() && link == p->right)
            {
                link = p;
                p = link->parent;
            }

            return p;
        }

        Link Predecessor(Link link)
        {
            if (link->left != Nil())
                return MaximumRBT(link->left);

            Link p = link->parent;
            while (p != Nil() && link == p->left)
            {
                link = p;
                p = link->parent;
            }

            return p;
        }

        int Heigth()
        {
            return HeightRBT(root);
        }

        int Length()
        {
            return LengthRBT(root);
        }

        void Clear()
        {
            ClearRBT(root);
        }

    private:
        void InsertRBT(Link& link, Link n)
        {
            if (link == RBT::Nil())
            {
                link = n;
                return;
            }

            if (n->item < link->item)
            {
                if (link->left == RBT::Nil())
                {
                    n->parent = link;
                    link ->left = n;
                }
                else
                    InsertRBT(link->left, n);
            }
            else
            {
                if (link->right == RBT::Nil())
                {
                    n->parent = link;
                    link->right = n;
                }
                else
                    InsertRBT(link->right, n);
            }

        }

        Link SearchRBT(Link link, Type item)
        {
            if (link == Nil() || link->item == item)
                return link;
            if (link->item > item)
                return SearchRBT(link->left, item);
            else
                return SearchRBT(link->right, item);
        }

        void InorderWalkRBT(Link link, std::ostream& out)
        {
            if (link == Nil())
                return;

            InorderWalkRBT(link->left, out);
            out << link->item << " ";
            InorderWalkRBT(link->right, out);
        }

        Link MaximumRBT(Link link)
        {
            while (link->right != Nil())
                link = link->right;
            return link;
        }

        Link MinimunRBT(Link link)
        {
            while (link->left != Nil())
                link = link->left;
            return link;
        }

        int HeightRBT(Link link)
        {
            if (link == Nil())
                return 0;

            int lh, rh;
            lh = HeightRBT(link->left);
            rh = HeightRBT(link->right);
            return std::max(lh, rh) + 1;
        }


        int LengthRBT(Link link)
        {
            if (link == Nil())
                return 0;
            return LengthRBT(link->left) + LengthRBT(link->right) + 1;
        }

        void ClearRBT(Link& link)
        {
            if (link == Nil())
                return;
            ClearRBT(link->left);
            ClearRBT(link->right);
            delete link;

            link = Nil();
        }

    private:
        Link root;
    };
}

#endif // RBT_H



测试代码如下：

#include <iostream>
#include <cassert>
#include <stdexcept>
#include "RBT.h"

using namespace std;
using namespace NTCI;

int main()
{
    try
    {
        int array[] = {5, 4, 7, 2, 6, 8, 9, 20, 12, 16, 11, 14};
        RBT<int> rbt(array, array + 12);

        cout << "排序后：";
        rbt.InorderWalk(cout);
        cout << endl;

        RBT<int>::Link link = rbt.Search(5);
        assert(rbt.Successor(link)->item == 6);
        assert(rbt.Predecessor(link)->item == 4);
        cout << "红黑树高度：" << rbt.Heigth() << endl;
        cout << "最大值：" << rbt.Maximum() << endl;
        cout << "最小值：" << rbt.Minimun() << endl;

        rbt.Remove(9);

        link = rbt.Search(7);
        assert(link->item == 7);

        link = rbt.Remove(link);
        delete link;
        assert(rbt.Length() == 10);
        cout << "删除元素7后：";
        rbt.InorderWalk(cout);
        cout << endl;

        rbt.Clear();
        assert(rbt.Length() == 0);

    }
    catch (exception& e)
    {
        cout << "发生异常：" << e.what() << endl;
    }
    return 0;
}