红黑树优化（C++实现）

前言

本文理论参考来源：

红黑树原理以及插入、删除算法附图例说明
关于AVL树和红黑树的一点看法
红黑树最多三次旋转达到平衡

由于红黑树的调整和旋转都是由下层至上层的，也就是说无论是左旋、右旋还是调整颜色,都需要当前结点的左右子树和父节点的地址，因此为了方便实现，红黑树结点的关键部分都被设计成以下的样子：

struct RBTreeNode {
    RBTreeNode *    lChild;
    RBTreeNode *    rChild;
    RBTreeNode *    parent;
    Color           color;
};

其大小一般在4*sizeof(void *)到3*sizeof(void *)不等，后者利用由于结点地址对齐到8B的原理，因此其成员指针底三位无法使用，可以用来存储颜色信息的性质将大小压缩到3个指针大小。

在下文，笔者尝试改进算法，将父节点指针取消掉，同时结合上面所说的利用地址对齐而不使用的低位比特存储颜色信息的方法，将红黑树结点的核心部分缩小至两个指针大小，即2*sizeof(void *)，实现插入、删除操作，在不改变时间复杂度的同时，将父节点的空间复杂度从O(n)减少为O(logn)。

思路

由于结点不携带父节点信息，所以需要在结点的遍历过程中发现并想办法存储各节点的父节点信息，在红黑树的插入、删除和遍历过程中，函数必然途径从根节点到目标结点的各层结点，因此可以想办法将其存储下来，以备后面使用。

思路1：递归

在每一层的结点向下层结点访问需要递归调用函数，每一层递归过程宏都需要分配一个自动指针变量指向对应层的结点。

以插入操作作为例子，插入操作需要在某一层或者若干层使用左右旋、调整颜色等，因此至少需要当前层的结点、父节点、祖父节点和曾祖父结点作为参数；同时，一旦发生旋转，上层各层对应的结点将会发生变化，因此上层传递到下层的指针参数必然是引用参数，这样下层旋转时可以顺带修改上层指针指向，上层操作不必重新调整自身指针，因此插入方法应该是类似下面这个样子的：

/**
* C:当前层结点
* P:上层结点
* G:上上层结点
* E:曾祖父结点
* T:目标结点
*/
void RBTree::insert(RBNode * & P,RBNode * & G,RBNode * & E,RBNode * T)
{
    //在当前递归层分配指向对应层结点的指针
    RBNode * C   = P->child;
    //....
    //由于向下层调用是传引用，因此下层可以修改C指针的指向
    insert(C,P,G,T);
    //哪怕下层发生了旋转，C指针发生了改变，C指针仍然指向该层对应结点
    //...
}

与下面的思路2相比，该思路不需要预知红黑树的层数，只要函数调用栈足够大，算法可以处理高度很高的红黑树，但是：

每一层递归都需要在栈帧中保存上层的5个参数即5*sizeof(void *)个字节，多层递归会占用大量的栈空间，而且递归调用必然额外添加函数的调用和返回等开销；

该思路涉及参数太多，给代码的编写与维护带来很大的麻烦，笔者在尝试编码的时候经常弄混，算法的复杂性决定这种递归调用最好是间接递归，即某一个递归函数可能在下层递归调用另一种递归函数，不同种类的递归函数之间就需要区分是否对应红黑树的同一层。

综上所述，本文不打算采用该思路。

思路2：迭代+自定义栈

顶层函数自动分配足够长的数组作为自定义栈，栈顶指向当前需要处理的结点。此思路主要的问题是多长的栈才是足够的长，可以在使用前预先确定红黑树最大的高度，如果红黑树最多容纳 $n$ 个结点,那么红黑树最大的高度为 $2\log_2(n)$ ；或者根据当前红黑树高度使用编译器扩展动态指定数组的长度。

实现

红黑树结点的定义

这里使用侵入式容器的思想，红黑树操作的结点仅为实际结点结构的一个字段，即“钩子”，实际结点的内存由使用者管理，红黑树仅负责根据“钩子”组织结点。

//rbtree.h

enum Color{REB = 0,BLACK = 1};

//红黑树钩子类
class RBHook
{
    //作为算法的内部数据，RBHook的所有方法和字段都不应该被暴露出来
    //只有算法和容器本身可以访问这些内容

    friend class RBTreeBase;

    template <typename A>
    friend class RBTree;

    uintptr_t_t data[2];

    RBHook * getLeftChild()const;

    RBHook * getRightChild()const;

    Color getColor()const;

    void setLeftChild(RBHook * L);

    void setRightChild(RBHook * R);

    void setColor(Color color);

    //清空指针，颜色设置为红色
    void clear();

    //与O指针指向的结点交换颜色
    void swapColor(RBHook * O);

} __attribute__((aligned(sizeof(void *))));

static_assert(sizeof(RBHook) == 2 * sizeof(void *),"RBHook大小不合规定");

关于“钩子”各方法的实现：

//rbtree.cpp
//data[0]高位存储左子节点地址、底1位存储颜色信息
//data[1]存储右子节点地址

RBHook * RBHook::getLeftChild()const
{
    return reinterpret_cast<RBHook *>(data[0] & (~ 1ULL));
}

RBHook * RBHook::getRightChild()const
{
    return reinterpret_cast<RBHook *>(data[1]);
}

Color RBHook::getColor()const
{
    return reinterpret_cast<Color>(data[0] & 1ULL);
}

void RBHook::setLeftChild(RBHook * pNode)
{
    data[0] = (data[0] & 1ULL) | reinterpret_cast<uintptr_t>(pNode);
}

void RBHook::setRightChild(RBHook * pNode)
{
    data[1] = reinterpret_cast<uintptr_t>(pNode);
}

void RBHook::setColor(Color color)
{
    data[0] = (color == BLACK) ? data[0] | 1ULL : data[0] & (~ 1ULL);
}

void RBHook::clear()
{
    data[0] = data[1] = 0;
}

void RBHook::swapColor(RBHook * O)
{
    Color color = O->getColor();
    O->setColor(getColor());
    setColor(color);
}

红黑树容器定义

为了使得容器具有泛用性，笔者将容器设计为泛型，实例化容器需要提供一个类型，该类型提供calc静态方法计算每一个结点的键值并指定键值的类型，键值类型需要有或者重载小于运算符，下面是一个实例化的例子：

//main.cpp

struct Data
{
    using ValueType = int;

    RBHook hook;
    int data;

    static ValueType calc(RBHook * H)
    {
        return reinterpret_cast<Data *>(H)->data;
    }
};

using Tree = RBTree<Data>;

同时红黑树算法中存在大量与泛型无关的代码，笔者又将这部分代码分离到单独的编译单元中，封装到红黑树基类RBTreeBase中：

//rbtree.h

//红黑树基类
//封装一切模板无关代码
class RBTreeBase
{
public:

    //初始化
    void init();

protected:

    //自定义栈类
    struct Stack
    {
        RBHook * *  pData;

        size_t      nSize;

        void push(RBHook * C);

        //弹出nLength个指针
        void pop(size_t nLength);

        //获取向栈底方向偏移栈顶nOffset的元素，如果nOffset超限，那么返回空指针
        RBHook * get(size_t nOffset);

        //设置向栈底方向偏移栈顶nOffset的元素
        void set(RBHook * C,size_t nOffset);

        //获取向栈顶方向偏移栈底nIdx的元素
        RBHook * fetch(size_t nIdx);

        //设置向栈顶方向偏移栈底nIdx的元素
        void update(RBHook * C,size_t nIdx);

        //返回栈顶相对栈底的偏移
        size_t topIdx();
    };

    //根节点
    RBHook *    pRoot;

    //结点数目
    uint32_t    nNumber;

    //当前红黑树可能的最大高度，会随着插入操作动态更新
    uint32_t    nMaxHeigh;


    //插入后的旋转调整，S的栈顶必须指向新插入的结点
    void adjustAfterInsert(Stack * S);

    //结点的替换、调整和删除，S的栈顶必须指向待删除的结点
    void removeImpl(Stack * S);

private:

    //结点的左旋
    //S栈顶为旋转结点的右子节点
    //旋转完成后，S高度减一，原旋转结点指针变为栈顶且指向原右子节点
    void leftRotate(Stack * S);

    //结点的右旋
    //S栈顶为旋转结点的左子节点
    //旋转完成后，S高度减一，原旋转结点指针变为栈顶且指向原左子节点
    void rightRotate(Stack * S);

};

最终容器的定义：

//rbtree.h

template <typename A>
class RBTree:public RBTreeBase
{

    using CalcType = A;

public:

    //结点的键值类型
    using ValueType = typename CalcType::ValueType;

    //插入操作总函数
    //成功返回true
    //重复插入、重复键值将导致插入失败并返回false
    bool insert(RBHook * T)
    {
        T->clear();
        if (pRoot == nullptr)
        {
            pRoot   = T;
            nNumber = nMaxHeigh = 1;
            T->setColor(BLACK);
            return true;
        }

        //下面是容器非空时插入操作
        //根据过程中可能的最大高度预先分配好栈空间
        //变长的自动数组需要gcc或者clang的扩展
        RBHook * stackData[nMaxHeigh + 1];
        Stack S{.pData = stackData,.nSize = 0};

        //下面将会将新节点作为叶子结点插入到树中
        //同时自定义栈会压入从新节点到根节点路径的所有结点
        S.push(pRoot);

        RBHook * C  = pRoot;
        while (true)
        {
            ValueType vC = CalcType::calc(C);
            ValueType vT = CalcType::calc(T);

            if (vC < vT)
            {
                if (C->getRightChild() == nullptr)
                {
                    C->setRightChild(T);
                    S.push(T);
                    break;
                }else
                {
                    C = C->getRightChild();
                    S.push(C);
                }
            }else if (vT < vC)
            {
                if (C->getLeftChild() == nullptr)
                {
                    C->setLeftChild(T);
                    S.push(T);
                    break;
                }else
                {
                    C = C->getLeftChild();
                    S.push(C);
                }
            }else return false;
        }

        //此时红黑树尚未调整，自定义栈的高度为红黑树在整个插入过程的最大高度
        nMaxHeigh   = nMaxHeigh > S.nSize ? nMaxHeigh : S.nSize;

        //调整红黑树
        adjustAfterInsert(&S);

        return true;
    }

    //删除操作总函数
    //成功返回true
    //如果容器空、待删除结点不在容器中、或者存在键值相同的不同结点，失败返回false
    bool remove(RBHook * T)
    {
        if (pRoot == nullptr) return false;

        if (nNumber == 1)
        {
            if (pRoot != T) return false;
            pRoot   = 0;
            nNumber = 0;
            return true;
        }

        RBHook * stackData[nMaxHeigh];
        Stack stack{.pData = stackData,.nSize = 0};
        Stack * S   = & stack;

        //下面根据红黑树结构找到待删除结点
        //同时自定义栈压入从待删除结点到根节点路径的所有节点
        S->push(pRoot);

        RBHook * C = pRoot;
        RBHook * P;
        ValueType vT = CalcType::calc(T);
        while (true)
        {
            ValueType vC    = CalcType::calc(C);
            if (vC < vT)
            {
                C   = C->getRightChild();
                if (C == nullptr) return false;
                S->push(C);
            }else if (vT < vC)
            {
                C   = C->getLeftChild();
                if (C == nullptr) return false;
                S->push(C);
            }else if (C != T) return false;
            else break;
        }
        
        //上面都是删除的准备工作
        //下面开始删除
        removeImpl(S);

        return true;
    }

};

算法实现——旋转

只介绍左旋，右旋类似

//rbtree.cpp

void RBTreeBase::leftRotate(Stack * S)
{
    //旋转结点的右子节点
    RBHook * C  = S->get(0);
    //旋转结点
    RBHook * P  = S->get(1);
    //旋转结点的父节点
    RBHook * G  = S->get(2);
    RBHook * L  = C->getLeftChild();

    //C、P、G、L从低层向上的顺序为：
    //L-C-P-G
    //除了L之外，其余三个结点按顺序在栈中

    if (G == nullptr) pRoot   = C;
    else if (G->getLeftChild() == P) G->setLeftChild(C);
    else G->setRightChild(C);

    P->setRightChild(L);
    C->setLeftChild(P);

    //此时C、P、G、L从低层向上的顺序为：
    //L-P-C-L
    //一般L和P不需要后续的操作，所以弹出一层
    //同时调整栈顶指向C，原来指向P
    S->pop(1);
    S->set(C,0);
}

算法实现——插入后调整

具体各种情况的操作描述见文章前言链接

//rbtree.cpp

void RBTreeBase::adjustAfterInsert(Stack * S)
{
    RBHook * C;
    //可能出现多次调色和旋转
    while (true)
    {
        //当前结点即栈顶
        C = S->get(0);
        //父节点
        RBHook * P = S->get(1);

        //当前结点为根节点
        if (P == nullptr)
        {
            C->setColor(BLACK);
            break;
        }

        //当前结点非根节点且父节点黑色
        if (P->getColor() == BLACK) break;

        //否则考察叔结点，此时祖父结点必存在
        //祖父结点
        RBHook * G = S->get(2);
        Bool bIsLP = G->getLeftChild() == P;
        //叔结点
        RBHook * W = bIsLP ? G->getRightChild() : G->getLeftChild();

        //叔结点非空且为红色
        if (W != nullptr && W->getColor() == RED)
        {
            W->setColor(BLACK);
            P->setColor(BLACK);
            G->setColor(RED);
            //需要变换当前结点为祖父结点，因此弹栈
            S->pop(2);
            //回到循环开头，重新分析
            continue;
        }

        //即左左、左右、右右、右左四种情况
        if (bIsLP)
        {
            if (P->getRightChild() == C) {
                leftRotate(S);
                P = C;
            }else S->pop(1);
            P->swapColor(G);
            rightRotate(S);
        }else
        {
            if (P->getLeftChild() == C) {
                rightRotate(S);
                P = C;
            }else S->pop(1);
            P->swapColor(G);
            leftRotate(S);
        }
        break;

    }

    //不要忘记更新数量
    nNumber += 1;
}

算法实现——替换、调整、删除