代码

#include <iostream>
#include <cstring>
#include <cstdio>
#include <cmath>
using namespace std;

template <typename T>
inline void read(T &k)
{
	int f = 1; k = 0; char c = getchar();
	while (c < '0' || c > '9') c == '-' && (f = -1), c = getchar();
	while (c <= '9' && c >= '0') k = k * 10 + c - '0', c = getchar();
	k *= f;
}

class RBTree
{
    static const int inf = 2 << 30 | 1;
private:
    enum ColorSet {Red, Black};

    struct Node
    {
        int Val, Size;
        Node *LT, *RT, *Fa;
        bool Col;

        Node(int _Val) : Val(_Val)
        {
            Size = 1;
            Col = Red;
            LT = RT = Fa = nullptr;
        }

        ~Node() {delete LT; delete RT;}

    };

    Node *Root;

public:
    RBTree() {Root = nullptr;}
    ~RBTree() {delete Root;}

    void Update(Node* const &x)
    {
        x->Size = (x->LT ? x->LT->Size : 0) + (x->RT ? x->RT->Size : 0) + 1;
    }

    void LeftRotate(Node* const &x)
    {
        Node *RT = x->RT;

        x->RT = RT->LT;
        if (RT->LT)
            RT->LT->Fa = x;

        RT->Fa = x->Fa;

        if (x->Fa)
            if (x->Fa->LT == x)
                x->Fa->LT = RT;
            else
                x->Fa->RT = RT;
        else
            Root = RT;

        x->Fa = RT;

        RT->LT = x;

        Update(x); Update(RT);
    }

    void RightRotate(Node* const &x)
    {
        Node *LT = x->LT;

        x->LT = LT->RT;
        if (LT->RT)
            LT->RT->Fa = x;

        LT->Fa = x->Fa;
        if (x->Fa)
            if (x->Fa->RT == x)
                x->Fa->RT = LT;
            else
                x->Fa->LT = LT;
        else
            Root = LT;

        x->Fa = LT;

        LT->RT = x;

        Update(x); Update(LT);
    }

    Node* Find(int Val, int ty = 0)
    {
        Node *x = Root, *Fa = nullptr;
        while (x)
        {
            Fa = x;
            if (!ty && x->Val == Val) return x;
            x = x->Val < Val ? x->RT : x-> LT;
        }

        return Fa;
    }

    void Insert(int Val)
    {
        Node *x = new Node(Val);

        Node* Fa = Find(Val, 1);
        x->Fa = Fa;
        if (Fa)
            if (Fa->Val < Val)
                Fa->RT = x;
            else
                Fa->LT = x;
        else
            Root = x;

        Node* y = Fa;
        while (y)
            ++y->Size, y = y->Fa;

        while (Fa && Fa->Col == Red)
        {
            Node* Gfa = Fa->Fa;
            Node* Unc = (Fa == Gfa->LT) ? Gfa->RT : Gfa->LT;

            if (Unc && Unc->Col == Red)
            {
                Fa->Col = Unc->Col = Black;
                Gfa->Col = Red;
                x = Gfa; Fa = x->Fa;
            }
            else
            {
                if (Fa == Gfa->LT)
                {
                    if (x == Fa->RT)
                        LeftRotate(Fa), swap(x, Fa);
                
                    Fa->Col = Black;
                    Gfa->Col = Red;
                    RightRotate(Gfa);
                }
                else
                {
                    if (x == Fa->LT)
                        RightRotate(Fa), swap(x, Fa);

                    Fa->Col = Black;
                    Gfa->Col = Red;
                    LeftRotate(Gfa);
                }
            }
        }

        if (!Fa && x->Col == Red)
            x->Col = Black;
    }

    void Erase(int Val)
    {
        Node *x = Find(Val);
        Node *Fa = x->Fa, *p = nullptr;
        bool DelCol = x->Col;

        if (!x->LT)
        {
            Node *RT = x->RT;
            p = RT;

            if (Fa)
                if (Fa->LT == x)
                    Fa->LT = RT;
                else
                    Fa->RT = RT;
            else
                Root = RT;

            if (RT)
                RT->Fa = Fa;
        }
        else if (!x->RT)
        {
            Node *LT = x->LT;
            p = LT;

            if (Fa)
                if (Fa->LT == x)
                    Fa->LT = LT;
                else
                    Fa->RT = LT;
            else
                Root = LT;

            if (LT)
                LT->Fa = Fa;
        }
        else
        {
            Node *y = x->RT;
            while (y->LT)
                y = y->LT;

            DelCol = y->Col;

            Node* RT = y->RT;
            p = RT;
            if (y->Fa != x)
            {
                Fa = y->Fa;
                Fa->LT = y->RT;
                if (y->RT)
                    y->RT->Fa = Fa;

                y->RT = x->RT;
                x->RT->Fa = y;
            }
            else
                Fa = y;

            if (x->Fa)
                if (x->Fa->LT == x)
                    x->Fa->LT = y;
                else
                    x->Fa->RT = y;
            else
                Root = y;

            y->Fa = x->Fa;
            y->LT = x->LT;
            x->LT->Fa = y;
            y->Col = x->Col;
            y->Size = x->Size;
        }

        Node *y = Fa;
        while (y)
            --y->Size, y = y->Fa;

        if (DelCol == Black)
        {
            x = p;
            while (x != Root && (!x || x->Col == Black))
            {
                if (x == Fa->LT)
                {
                    Node *Bro = Fa->RT;
                    if (Bro->Col == Red)
                    {
                        Bro->Col = Black;
                        Fa->Col = Red;
                        LeftRotate(Fa);
                        Bro = Fa->RT;
                    }

                    if ((!Bro->LT || Bro->LT->Col == Black) && (!Bro->RT || Bro->RT->Col == Black))
                    {
                        Bro->Col = Red;
                        x = Fa;
                        Fa = x->Fa;
                    }
                    else
                    {
                        if (!Bro->RT || Bro->RT->Col == Black)
                        {
                            Node *Nie = Bro->LT;
                            Nie->Col = Black;
                            Bro->Col = Red;
                            RightRotate(Bro);
                            Bro = Nie;
                        }

                        Bro->Col = Fa->Col;
                        Fa->Col = Black;
                        Bro->RT->Col = Black;
                        LeftRotate(Fa);
                        x = Root;
                    }
                }
                else
                {
                    Node* Bro = Fa->LT;
                    if (Bro->Col == Red)
                    {
                        Bro->Col = Black;
                        Fa->Col = Red;
                        RightRotate(Fa);
                        Bro = Fa->LT;
                    }
                    if ((!Bro->LT || Bro->LT->Col == Black) && (!Bro->RT || Bro->RT->Col == Black))
                    {
                        Bro->Col = Red;
                        x = Fa;
                        Fa = x->Fa;
                    }
                    else
                    {
                        if (!Bro->LT || Bro->LT->Col == Black)
                        {
                            Node* Nie = Bro->RT;
                            Nie->Col = Black;
                            Bro->Col = Red;
                            LeftRotate(Bro);
                            Bro = Nie;
                        }

                        Bro->Col = Fa->Col;
                        Fa->Col = Black;
                        Bro->LT->Col = Black;
                        RightRotate(Fa);
                        x = Root;
                    }
                }
            }

            if (x)
                x->Col = Black;
        }
    }

    int Get_Rank(int Val)
    {
        Node *x = Root;
        int Rank = 1; 
        while (x)
        {
            if (Val > x->Val)
                Rank += x->LT ? x->LT->Size + 1 : 1, x = x->RT;
            else
                x = x->LT;
        }

        return Rank;
    }

    int Get_Kth(int K)
    {
        Node *x = Root;
        while (x)
        {
            int LSize = x->LT ? x->LT->Size : 0;
            if (K == LSize + 1)
                return x->Val;

            if (K <= LSize)
                x = x->LT;
            else
                K -= LSize + 1, x = x->RT;
        }

        return inf;
    }

    int Get_Pre(int Val)
    {
        Node* x = Root;
        int ans = 0;
        while (x)
        {
            if (Val > x->Val)
                ans = x->Val, x = x->RT;
            else
                x = x->LT;
        }

        return ans;
    }

    int Get_Next(int Val)
    {
        Node *x = Root;
        int ans = 0;
        while (x)
        {
            if (Val < x->Val)
                ans = x->Val, x = x->LT;
            else
                x = x->RT;
        }

        return ans;
    }
};


int main()
{
    int n, ty, x;

    // read(n);
    n = 500000;
    int a[n + 1];
    RBTree Tr;
    for (int i = 1; i <= n; i++)
        Tr.Insert(a[i] = rand());
    for (int i = 1; i <= n; i++)
        Tr.Erase(a[i]);
    // for (int i = 1; i <= n; i++)
    // {
    //     read(ty); read(x);
    //     if (ty == 1) Tr.Insert(x);
    //     else if (ty == 2) Tr.Erase(x);
    //     else if (ty == 3) printf("%d\n", Tr.Get_Rank(x));
    //     else if (ty == 4) printf("%d\n", Tr.Get_Kth(x));
    //     else if (ty == 5) printf("%d\n", Tr.Get_Pre(x));
    //     else printf("%d\n", Tr.Get_Next(x));
    // }
}

评测结果