红黑树 Red Black Tree C++

@vrqq  July 27, 2020

欠的债早晚得还,之前总觉得太长记不住,总用Treap等奇技淫巧对付。。
下一个债会不会是数论呢。。。

序 红黑树规则

推荐先会写些简单的树例如splay / Treap 之后,再来看红黑树,好理解很多。
参考一个Treap树:https://blog.vrqq.org/archives/66/

红黑树就是满足下列条件的树

  • root根节点 是黑色的
  • 红点如果有子节点,那一定是黑的(也就是说红点的父节点也一定是黑点)
  • 新插入的点是红点,和二叉树惯用插入一样,放到现有leaf底下
  • 对于树上任何一个点,沿着任意一条路往下走 直到leaf点,路径上的黑点数都相同。

插入

参考 https://zhuanlan.zhihu.com/p/87469768
文中前面的图很清晰,简单来说就是:

  • 先找到合适的原有leaf,把“新点”放在原有leaf下面,当作新的leaf
  • 然后把“新点”,和“新点父亲”,“新点爷爷”,“新点伯伯”组成一个“八”字,然后旋转 or/and 标色,使之满足最后一条“路径上黑点数量相同”,和“红点的孩子是黑点”

无论是什么操作,依赖的准则都是:使用旋转,重新标色等方法,满足“变动过后的子树”,能融入原有“树林”中。

提问:为什么不使用移动,例如把一支直接安到另一个点上,也满足红黑树要求?
因为旋转可以带来左右分支高度差不多,如果直接移动“一个树杈”,则会造成某一“枝”极长。

删除

简单说下删除的过程,有个大概思路,然后推理一下,可以先想一下 再看参考文档。
(假设没有重复点)下文以删除“树中间的点”为例:

  • 先找到待删除的点(记为O),假设 O点同时有左子树和右子树。
  • 然后在其左子树/右子树里找到能替代这个位置的点,我们一般选“右子树里最左边的值”,也就是中序遍历(左根右)里的下一个点(记为N)。(下文例子就是这样)
  • 然后不改变树的结构和点的颜色,而是O.value = N.value,这样就出现了两个重复数值的点。
  • 然后删掉N点就ok了,N点虽然没有左子树,但可能还有右子树!

    • 如果N点是红点,直接删掉,把后代(右子树if have)接上就行,一定是黑点 不会犯冲;
    • 如果N点是黑点,那么删掉N点后,后代想接上就要面临:1)后代还是黑点;2)这一支N的后代,和森林里其他支比,欠一个黑点。
    • 如果N是黑点,后代是红点,那么把后代红点直接变黑就满足条件了

于是我们约定如下删除准则
设N点的父节点为P,叔叔节点为S (if have),叔叔的两个孩子是SL和SR (if have)
使用旋转,和重新标定颜色两种方法,使当移除N后,N的后代子树 连接到P之后,依然保持红黑树性质。
(“N的后代子树”是N的右子树,连接以后 “N的后代子树”是P的左子树)

解密时间

  • 目标是想给“N的后代那一支”多分一个黑色
  • 我们在尽量小的范围内动一动,所以我们只动“P子树”,而不关心森林里其他的分支;
  • 在P点上只使用Counter-Clockwise,也就是:N的父节点始终是P,“N的后代”接替N时候也连接在P的左支;
  • N后代的颜色” 一定是黑色;

也就是说:通过在“P子树”这一范围内闪转腾挪,以 向“S子树借 黑点”、重新标色等手段,使在删除N点,并将其后代连上后,新的“P子树”依然能融入大森林。

旋转wiki里面是2个手法:(下文是其中一组合理方式,理论上有也其他make sense方式,不知道为什么没有采用这里TBD)

  • 如果说需要旋转才能make sense,那么。。
  • 把"S的内侧孩子"向外旋转到 S的位置,然后再在P点 左旋/右旋
  • (如果是N,P,S组成八字)在P点 左旋/右旋;
  • 旋转之后重新标色

howtos?
需要准备很多纸,先把能想到的情况画一画,穷举 当面临某情况时候应该怎么处理。。(可以先不立刻分类)
<到这里时,鼓励自己想变换方式 而不是参考wiki>
然后参考下下述链接 Wikipedia "删除" 的分类,重新思考看能否覆盖刚才自己想想象的情况!(important)
最后根据Wiki里面列举5种情况,先在纸上把转换形式都画出来,然后“五种变换”之间如何衔接 用其他笔画个箭头,就明了了。
https://zh.wikipedia.org/wiki/%E7%BA%A2%E9%BB%91%E6%A0%91

以及我粗略的看了下 https://zhuanlan.zhihu.com/p/25402654

测试题
插入,查找count larger than 'key'
codeforces 1042d https://codeforces.com/problemset/problem/1042/D

删除
codeforces 375d https://codeforces.com/contest/375/problem/D

多一点思考

许多题可以用二叉树,但不能用树状数组?为啥,因为数据范围太大。
那么我们如果已知了所有可能出现的结果呢?离散化一下就可以了!
比如上述1042D,我们已知所有可能出现的值,那么把sum[i]离散化,然后再插入树状数组。搜索时多一次二分查找而已!

扩展:实现迭代器

以正向迭代器为例,我们需要达成如下两个:

  • std::prev(tree.end()) == iterator(tree.right_most_one())
  • ++iterator(tree.right_most_one()) == tree.end()

如果仅仅用iterator(nullptr) 表示end(), 就不能从end往前走。
后来从GCC STL代码发现了技巧:树的真正根节点用root表示,增设一虚拟节点header,动态维护如下关系:

  • header.left = tree.right_most_one()
  • tree.root.parent = header
  • tree.right_most_one().right == nullptr 这个不变

迭代器的加减是树的中序遍历(左根右),"right_most_one的后一个" 通过root.parent就找到了header,同样 "header(end())的前一个" 通过header.left就找到了树中的最右节点(right_most_one)。

Sample code

// Red-black tree
// vrqq (vrqq.org)
// 88240534  Jul/29/2020 07:31UTC+8  vrqq    D - Petya and Array GNU C++17   Accepted    155 ms  9900 KB
// 88426725 Jul/30/2020 18:09UTC+8  vrqq    D - Tree and Queries    GNU C++17   Accepted    623 ms  37300 KB
#include <iostream>
#include <sstream>
#include <vector>

template<typename T, int MAXN = -1>
class RBTree {
public:
    enum Color: bool {
        RED = 1,
        BLACK = 0
    };

    struct Node {
        T key;
        size_t size; // The size of subtree
        bool color;  // R==1 and B==0
        Node *left, *right, *parent;
        Node(T key, Node *parent) : key(key), parent(parent) {
            size = 1;
            color = RED;
            left = right = nullptr;
        }
        Node() {
            size = 1;
            color = RED;
            left = right = nullptr;
        }
        Node &operator=(const Node &rhs) {
            key = rhs.key;
            size = rhs.size;
            color = rhs.color;
            left = rhs.left;
            right = rhs.right;
            parent = rhs.parent;
        }
    }*root = nullptr;
    
private:

    void swapNode(Node *a, Node *b) {
        // std::swap(a->key, b->key);
        std::swap(a->size, b->size);
        std::swap(a->color, b->color);
        
        if (a->left)
            a->left->parent = b;
        if (b->left)
            b->left->parent = a;
        std::swap(a->left, b->left);

        if (a->right)
            a->right->parent = b;
        if (b->right)
            b->right->parent = a;
        std::swap(a->right, b->right);

        if (a->parent) //a and b may have same parent
            a->parent->left==a ? (a->parent->left = b) : (a->parent->right = b);
        if (b->parent)
            b->parent->left==b ? (b->parent->left = a) : (b->parent->right = a);
        std::swap(a->parent, b->parent);
    }

    template<int, class = void>
    struct storage;

    template<int NN>
    struct storage<NN, std::enable_if_t<(MAXN > 0)>> {
        Node *pool; //gcc do not support 'new Node[NN]' here.
        int nptr=0;
        std::vector<Node*> recycle;
        storage() {pool = new Node[NN];}
        Node *alloc(T key, Node *parent) {
            if (recycle.size()) {
                Node *rv = new (recycle.back()) Node{key, parent};
                recycle.pop_back();
                return rv;
            }
            if (nptr >= NN)
                throw std::bad_alloc();
            // std::cout<<"Preallocated"<<std::endl;
            return new (&pool[nptr++]) Node{key, parent};
        }
        void free(Node *p) {
            recycle.push_back(p);
        }
        ~storage() {
            delete[] pool;
        }
    };

    template<int NN>
    struct storage<NN, std::enable_if_t<(NN <= 0)>> {
        Node *alloc(T key, Node *parent) {
            // std::cout<<"New"<<std::endl;
            return new Node{key, parent};
        }
        void free(Node *p) {
            delete p;
        }
    };

    storage<MAXN> allocator;

    //Clockwise
    //The rotate must be inplaced since p have parent information
    Node *cw(Node *p);

    //Counter-clockwise
    //The rotate must be inplaced since p have parent information
    Node* ccw(Node *p);

    size_t count_larger_than(Node *p, const T& key) const;

    size_t count_smaller_than(Node *p, const T& key) const;

    void print_tree(Node *p);

    void askBlack(Node *p);

    void clear(Node *p) {
        if (p) clear(p->left), clear(p->right), allocator.free(p);
    }

    template<int NN>
    typename std::enable_if_t<(NN > 0)> destroy(Node *p) {}

    template<int NN>
    typename std::enable_if_t<(NN <= 0)> destroy(Node *p) {
        clear(p);
    }

public:
    RBTree() = default;
    RBTree(RBTree &&rhs) = delete;

    ~RBTree() {
        destroy<MAXN>(root);
    }

    size_t size() const { return root?root->size:0; }

    Node* insert(T newkey);

    size_t count_larger_than(const T& key) const { return count_larger_than(root, key); }

    size_t count_smaller_than(const T& key) const { return count_smaller_than(root, key); }

    void erase(Node *p);

    void clear() { clear(root); root = nullptr; }

    void debug() { print_tree(root); }
}; //END class RBTree


template<typename T, int MAXN>
typename RBTree<T, MAXN>::Node* 
RBTree<T, MAXN>::cw(Node *p) {
    Node *tmp = p->left;
    Node *up = p->parent;
    if (up) {
        if (up->left == p)
            up->left = tmp;
        else
            up->right = tmp;
    }
    tmp->parent = up;
    p->parent = tmp;
    p->left = tmp->right;
    if (tmp->right)
        tmp->right->parent = p;

    tmp->right = p;
    p->size = (p->left?p->left->size:0) + (p->right?p->right->size:0) + 1;
    tmp->size = (tmp->left?tmp->left->size:0) + p->size + 1;
    return tmp;
}


template<typename T, int MAXN>
typename RBTree<T, MAXN>::Node* 
RBTree<T, MAXN>::ccw(Node *p) {
        Node *tmp = p->right;
        p->right = tmp->left;
        if (p->right)
            p->right->parent = p;
        tmp->left = p;
        if (p->parent)
            p->parent->left==p? (p->parent->left=tmp): (p->parent->right=tmp);
        tmp->parent = p->parent;
        p->parent = tmp;
        p->size = (p->left?p->left->size:0) + (p->right?p->right->size:0) + 1;
        tmp->size = p->size + (tmp->right?tmp->right->size:0) + 1;
        return tmp;
    }


template<typename T, int MAXN>
typename RBTree<T, MAXN>::Node* RBTree<T, MAXN>::insert(T newkey) {
    // std::cout<<"++ Insert "<<newkey<<std::endl;
    //Insert
    Node *up = nullptr, *p = root, **up2down, *rv;
    while(p != nullptr) {
        p->size++;
        up = p;
        if (newkey < p->key)
            up2down = &p->left, p = p->left;
        else
            up2down = &p->right, p = p->right;
    }
    // p = new Node(newkey, up);
    p = allocator.alloc(newkey, up);
    if (up == nullptr)//special for null tree
        root = p;
    else
        *up2down = p;
    rv = p;

    // std::cout<<"\t append after node "<<(up?up->key:-1)<<std::endl;

    //maintain : p is the active node
    while(p != root && up != root && p->color == RED) {
        if (up->color == BLACK)
            break;
        //Now: both up and p are all RED, top is BLACK
        Node *top = up->parent;
        if (top->left == up) {
            if (top->right && top->right->color == RED) {
                up->color = top->right->color = BLACK;
                top->color = RED;
                p = top; up = p->parent; continue;
            }
            //the other situations below can be processed in 'Node p'-tree
            if (p == up->right) {
                up = ccw(up);
                p = up->left;
            }//Now the tree will be like '/\'
            
            std::swap(top->color, up->color);
            (top==root?root:top) = cw(top);
            break;
        }
        else { // top->right == up
            if (top->left == nullptr || top->left->color == BLACK) {
                if (p == up->left) {
                    up = cw(up);
                    p = up->right;
                }
                std::swap(up->color, top->color);
                (top==root?root:top) = ccw(top);
                break;
            }
            else {
                top->color = RED;
                top->left->color = top->right->color = BLACK;
                p = top; up = p->parent;
            }
        }
    }

    root->color = BLACK;
    return rv;
}


//Ask an extra BLACK-node for the 'p-subtree'
//* for any, ? for optional (can be nullptr)
//Example for 'p' is the left branch of 'up'.
//Case 1: up=B, sibling=R, sibling.child=[B?, B?]  => Case 2, 3, 4
//Case 2: up=R, sibling=B, sibling.child=[B?, B?]  => >>return
//Case 3: up=*, sibling=B, sibling.child=[R, B?]   => Case 4
//Case 4: up=*, sibling=B, sibling.child=[*?, R]   => >>return
//Case 5: up=B, sibling=B, sibling.child=[B?, B?]  => recursion(up)
template<typename T, int MAXN>
void RBTree<T,MAXN>::askBlack(Node *p) {
    // std::cout<<"  :: ASK for BLACK: "<<p->key<<std::endl;
    if (p == root) {
        p->color = BLACK;
        return ;
    }
    Node *up = p->parent;
    Node *sibling = (p==up->left ? up->right : up->left);
    bool at_left = (p == up->left);

    if (sibling == nullptr) {
        if (up->color == BLACK)
            return askBlack(up);
        up->color = BLACK;
        return ;
    }

    //Case 1: up=B, sibling=R, sibling.child=[?, ?]
    if (up->color == BLACK && sibling->color == RED) {
        // std::cout<<"  ::-> Case 1"<<std::endl;
        at_left?ccw(up):cw(up);
        if (root == up)
            root = up->parent;
        std::swap(up->color, up->parent->color);
        sibling = (at_left ? up->right : up->left);
    }

    //Case 2: up=R, sibling=B, sibling.child=[B?, B?]
    if (up->color == RED && sibling->color == BLACK
        && (sibling->left?sibling->left->color:BLACK) == BLACK
        && (sibling->right?sibling->right->color:BLACK) == BLACK) {
            // std::cout<<"  ::-> Case 2"<<std::endl;
            std::swap(sibling->color, up->color);
            return ;
        }

    //Case 3: up=*, sibling=B, sibling.child=[R, B?] when up is the left of p-childs
    if (sibling->color == BLACK) {
        Node *s_inner = at_left?sibling->left:sibling->right;
        Node *s_outer = at_left?sibling->right:sibling->left;
        if (s_inner && s_inner->color == RED && (s_outer?s_outer->color==BLACK:true)) {
            // std::cout<<"  ::-> Case 3"<<std::endl;
            std::swap(sibling->color, s_inner->color);
            sibling = at_left?cw(sibling):ccw(sibling);
        }
    }

    //Case 4: up=*, sibling=B, sibling.child=[*?, R] when up is the left of p-childs
    if (sibling->color == BLACK) {
        Node *s_outer = at_left?sibling->right:sibling->left;
        if (s_outer && s_outer->color == RED) {
            // std::cout<<"  ::-> Case 4"<<std::endl;
            std::swap(sibling->color, up->color);
            s_outer->color = BLACK;
            at_left?ccw(up):cw(up);
            if (root == up)
                root = up->parent;
            return ;
        }
    }

    //Case 5: up=B, sibling=B, sibling.child=[B?, B?]
    if (up->color == BLACK && sibling->color == BLACK &&
        (sibling->left?sibling->left->color==BLACK:true) &&
        (sibling->right?sibling->right->color==BLACK:true)) {
            // std::cout<<"  ::-> Case 5"<<std::endl;
            sibling->color = RED;
            return askBlack(up);
        }
}


template<typename T, int MAXN>
void RBTree<T,MAXN>::erase(Node *p) {
    if (p->left && p->right) {
        Node *nxt = p->right;
        while(nxt->left)
            nxt = nxt->left;
        // std::cout<<"\nSelect pending: "<<nxt->key<<" after swap "<<std::endl;
        swapNode(p, nxt);
        if (root == p)
            root = nxt;
    }

    //Now: p has only one or zero child.
    Node* progeny = p->left?p->left:(p->right?p->right:nullptr);
    if (p->color == BLACK) {
        if (progeny && progeny->color == RED)
            progeny->color = BLACK;
        else
            askBlack(p);
    }

    if (progeny)
        progeny->parent = p->parent;

    if (p->parent) {
        if (p->parent->left == p)
            p->parent->left = progeny;
        else
            p->parent->right = progeny;
        for (auto t = p->parent; t; t=t->parent)
            t->size--;
    }
    else
        root = progeny;
    allocator.free(p);
}


template<typename T, int MAXN>
void RBTree<T, MAXN>::print_tree(Node *p) {
    if (!p)
        return ;
    auto fmt = [&](Node *p) {
        std::stringstream ss;
        if (p == nullptr)
            ss<<"NIL";
        else
            ss<<p->key;
        return ss.str();
    };
    std::cout<<"-- " << p->key << (p->color==RED?" RED  ":" BLACK")
        << " >> child=["
        <<fmt(p->left)<<", "
        <<fmt(p->right)<<"]"
        <<"\tup=" << fmt(p->parent)
        <<" size=" <<p->size <<std::endl;
    print_tree(p->left);
    print_tree(p->right);
}


template<typename T, int MAXN>
size_t RBTree<T, MAXN>::count_larger_than(Node *p, const T& key) const {
    if (p == nullptr)
        return 0;
    if (!(key < p->key)) // p->key <= key
        return count_larger_than(p->right, key);
    return count_larger_than(p->left, key) + (p->right?p->right->size:0) + 1;
}


template<typename T, int MAXN>
size_t RBTree<T, MAXN>::count_smaller_than(Node *p, const T& key) const {
    if (p == nullptr)
        return 0;
    if (!(p->key < key)) // key <= p->key
        return count_smaller_than(p->left, key);
    return count_smaller_than(p->right, key) + (p->left?p->left->size:0) + 1;
}

添加新评论