LibreOJ #107. 维护全序集

/*
    LibreOJ #107. 维护全序集
    Splay
    支持 
    1.把 x 加入 S
    2.删除 S 中的一个 x，保证删除的 x 一定存在
    3.求 S 中第 k 小
    4.求 S 中有多少个元素小于 x
    5.求 S 中小于 x 的最大数
    6.求 S 中大于 x 的最小数

    我之前用Splay找数x的排名时采用的是挨个找的
    现在发现根本没必要
    只需要把splay(x)后减去右子树的大小与根节点的数出现几次即可 
　　许久没写splay
　　写的比较丑。。

*/
#include <cstdio>

void read (int &now)
{
    register char word = getchar ();
    bool temp = false;
    for (now = 0; word < '0' || word > '9'; word = getchar ())
        if (word == '-')
            temp = true;
    for (; word >= '0' && word <= '9'; now = now * 10 + word - '0', word = getchar ());
    if (temp)
        now = -now;
}


struct S_D 
{
    
    S_D *child[2], *father;
    
    int value;
    int weigth, size;
    
    S_D () 
    {
        child[0] = child[1] = NULL;
        father = NULL;
        value = 0;
        size = 1;
        weigth = 1;
    }
    
    inline int Get_Pos ()
    {
        return this->father->child[1] == this;
    }
    
    inline void Updata ()
    {
        this->size = this->weigth;
        if (this->child[0])
            this->size += this->child[0]->size;
        if (this->child[1])
            this->size += this->child[1]->size;
    }
    
    inline void Clear ()
    {
        child[0] = child[1] = NULL;
        father = NULL;
        value = size = weigth = 0;
    }
};

class Splay_Tree_Type
{
    
    private :
        
        S_D *Root;

        inline void Rotate (S_D *now)
        {
            S_D *Father = now->father;
            register int pos = now->Get_Pos () ^ 1;
            Father->child[pos ^ 1] = now->child[pos];
            if (now->child[pos])
                now->child[pos]->father = Father;
            if ((now->father = Father->father) != NULL)
                now->father->child[now->father->child[1] == Father] = now;
            Father->father = now;
            now->child[pos] = Father;

            Father->Updata ();
            now->Updata ();
            
            if (now->father == NULL)
                Root = now;
        }
        
        inline void Splay (S_D *now)
        {
            for (S_D *Father; Father = now->father; Rotate (now))
                if (Father->father)
                    Rotate (now->Get_Pos () == Father->Get_Pos () ? Father : now);
        }
        
    public :
        
        void Insert (int key)
        {
            if (Root == NULL)
            {
                Root = new S_D;
                Root->value = key;
                return ;
            }
            S_D *now = Root, *Father;
            for (; ;Father = now, now = now->child[key > now->value])
            {
                if (now == NULL)
                {
                    now = new S_D;
                    now->father = Father;
                    now->value = key;
                    Father->child[key > Father->value] = now;
                    Splay (now);
                    return ;
                }
                if (now->value == key)
                {
                    now->weigth ++;
                    Splay (now);
                    return ;
                }
            }
        }
        
        S_D *Find (int key)
        {
            S_D *now = Root;
            for (S_D *Father; (now != NULL && now->value != key); Father = now, now = now->child[key > now->value]);
            if (now == NULL)
                return now;
            Splay (now);
            return now;
        }
        
        int Find_Prifix (int key)
        {
            S_D *now = Find (key);
            bool flag = false;
            if (now == NULL)
            {
                this->Insert (key);
                now = Root;
                flag = true;
            }
            if (now->child[0] == NULL)
            {
                if (flag)
                    this->Delete (key);
                return -1;
            }
            for (now = now->child[0]; now->child[1]; now = now->child[1]);
            if (flag)
                this->Delete (key);
            return now->value;
        }
        
        S_D* Find_Prifix_pos (int key)
        {
            S_D *now = Find (key);
            bool flag = false;
            if (now == NULL)
            {
                this->Insert (key);
                now = Root;
                flag = true;
            }
            if (now->child[0] == NULL)
            {
                if (flag)
                    this->Delete (key);
                return NULL;
            }
            for (now = now->child[0]; now->child[1]; now = now->child[1]);
            if (flag)    
                this->Delete (key);
            return now;
        }
        
        int Find_Suffix (int key)
        {
            S_D *now = Find (key);
            bool flag = false;
            if (now == NULL)
            {
                this->Insert (key);
                now = Root;
                flag = true;
            }
            if (now->child[1] == NULL)
            {
                if (flag)
                    this->Delete (key);
                return -1;
            }
            for (now = now->child[1]; now->child[0]; now = now->child[0]);
            if (flag)
                this->Delete (key);
            return now->value;
        }
        
        int Get_rank (int key)
        {
            S_D *now = Find (key);
            bool flag = false;
            if (now == NULL)
            {
                this->Insert (key);
                flag = true;
            }
            int Answer = Root->size;
            if (Root->child[1])
                Answer -= Root->child[1]->size;
            Answer -= Root->weigth;
            if (flag)
                this->Delete (key);
            return Answer;
        }
        
        void Delete (int key)
        {
            S_D *now = Find (key);
            if (now->weigth > 1)
            {
                now->weigth --;
                now->size --;
                return ;
            }
            if (now->child[0] == NULL && now->child[1] == NULL)
            {
                now->Clear ();
                now = NULL; // ????could change to delete
                Root = NULL;
                return ;
            }
            if (now->child[0] == NULL)
            {
                S_D *res = now;
                Root = now->child[1];
                now->child[1]->father = NULL;
                now->Clear ();
                now = NULL;
                return ;
            }
            if (now->child[1] == NULL)
            {
                S_D *res = now;
                Root = now->child[0];
                now->child[0]->father = NULL;
                now->Clear ();
                now = NULL;
                return ;
            }
            
            S_D *res_pre = Find_Prifix_pos (now->value);
            S_D *res = Root;
            Splay (res_pre);
            
            Root->child[1] = res->child[1];
            res->child[1]->father = Root;
            res->Clear ();
            Root->Updata ();
        }
        
        int Get_kth_number (int key)
        {
            int Answer = 0;
            register int res;
            for (S_D *now = Root; ;)
            {
                if (now->child[0] && key <= now->child[0]->size)
                {
                    now = now->child[0];
                    continue;
                }
                res = (now->child[0] ? now->child[0]->size : 0) + now->weigth;
                if (key <= res)
                    return now->value;
                key -= res;
                now = now->child[1];
            }
        }
};

Splay_Tree_Type Splay;

int main (int argc, char *argv[])
{
    int N;
     read (N);
    int type, x;
    
    for (register int i = 1; i <= N; i ++)
    {
        read (type);
        read (x);
        switch (type)
        {
            case 0: 
            {
                Splay.Insert (x);
                break;
            }
            case 1:
            {
                Splay.Delete (x);
                break;
            }
            case 2:
            {
                printf ("%d
", Splay.Get_kth_number (x));
                break;
            }
            case 3:
            {
                printf ("%d
", Splay.Get_rank (x));
                break;
            }
            case 4:
            {
                printf ("%d
", Splay.Find_Prifix (x)); 
                break;
            }
            case 5:
            {
                printf ("%d
", Splay.Find_Suffix (x)); 
                break;
            }
        }
    }
    
    return 0;
}