PBDS 구현하기

일반적으로 STL의 std::set은 k-th element를 찾거나 counting을 log 시간에 할 수 없다.

PBDS(Policy Based Data Structure)은 이를 해결할 수 있다고 한다.

 

Set과 같은 정렬 형태의 자료구조가 필요하면서 위의 동작 요구하는 BOJ의 일반적인 문제는 입력의 제한이 있기 때문에 세그먼트 트리 등으로 충분히 해결할 수 있는 경우가 대부분이다.

다만, interaction 문제나 세그먼트 트리로 만들어야하는 key 값의 range가 매우 크거나 실시간 쿼리가 필요한 경우 PBDS를 통해 해결하는 것이 조금 더 효율적일 것이다.

 

사실 BBST(Balanced Binary Search Tree)라면 PBDS를 쉽게 구현할 수 있는 것 같아서 유명한 RB Tree(Red-Black Tree)를 활용하여 구현하였다.

* PBDS의 실제 코드를 참고하지는 않아서 최적화나 정확성이 떨어지는 코드일 수 있다.

 

IsK(int k) : key에 맞는 Node 포이터를 구할 수 있다

IsN(int n) : k-th element 값을 구할 수 있다

 

key value가 int형이 아니라면 Node 구조체에서 수정하고 비교연산자 오버라이딩을 하면 될 것 같다

 

아래는 BOJ 2243번을 해결하는 소스코드이다.

 

#include <bits/stdc++.h>
using namespace std;
#define endl '\n'
typedef long double ld;
typedef long long ll;
typedef __int128_t lll;
typedef pair<ll, ll> pll;

enum Color{Red,Black};

struct Node{
    int k=0; // key
    int n=0, ln=0, rn=0; // num
    Node *p=nullptr, *l=nullptr, *r=nullptr; // parent, left, right
    Color c = Black; // color
};

class RBTree{
public: // private
    int n;
    Node *root, *nil;

    pair<Node*, Node*> IsK(int k){
        Node *t = root, *p = nullptr;
        for(; t!=nil && t->k != k; p = t, t = (k < p->k) ? p->l : p->r);
        return {t, p};
    }
    int IsN(int n){
        for(Node *t = root; t != nil;){
            if(n <= t->ln) t=t->l;
            else if(n <= t->ln+t->n){
                int ret = t->k;
                Delete(t->k, 1);
                return ret;
            }
            else n -= t->ln+t->n, t=t->r;
        }
    }
    void Insert(int k, int n){
        Node *z;
        auto[x,y] = IsK(k);
        if(x == nil) { // no key k
            z = new Node{k,n,0,0,y,nil,nil,Red};
            if(y == nullptr) root = z, z->c = Black;
            else {
                if(z->k > y->k) y->r = z;
                else y->l = z;
                if(y->p != nullptr) InsertFix(z);
            }
        }
        else z = x, z->n += n;
        for(Node *p = z->p; p!=nullptr; z=p, p=p->p) (p->l == z) ? p->ln += n : p->rn += n;
        return;
    }
    void InsertFix(Node *z){
        while(z!=root && z->p->c == Red){
            Node *q = z->p->p; // grand parent
            Node *u = (z->p == q->l) ? q->r : q->l; // uncle
            bool  s = (z->p == q->l);

            if(u && u->c == Red){
                z->p->c = u->c = Black, q->c = Red;
                z = q;
            }
            else{
                if(z==(s?z->p->r:z->p->l)) z=z->p, s?RotL(z):RotR(z);
                z->p->c = Black, q->c = Red;
                s?RotR(q):RotL(q);
            }
        }
        root->c = Black;
    }
    void Delete(int k, int n){
        auto[x,y] = IsK(k);
        if(x==nil) return;
        Node *z = x;
        for(z->n -= n; y!=nullptr; x=y, y=x->p) (y->l == x) ? y->ln -= n : y->rn -= n;
        if(!z->n){
            Color c = z->c;
            if(z->l == nil)      x = z->r, Trans(z, z->r);
            else if(z->r == nil) x = z->l, Trans(z, z->l);
            else{
                y = tree_minimum(z->r);
                c = y->c;
                x = y->r;

                for(Node *t=y; t->p!=z; t=t->p) (t->p->l == t) ? t->p->ln -= y->n : t->p->rn -= y->n;

                if(y->p == z) x->p = y;
                else{
                    Trans(y, y->r);
                    y->r = z->r;
                    y->r->p = y;
                }
                Trans(z, y);
                y->l = z->l;
                y->l->p = y;
                y->ln = y->l->ln + y->l->n + y->l->rn;
                y->rn = y->r->ln + y->r->n + y->r->rn;
                y->c = z->c;
            }
            delete z;
            if(c == Black) DeleteFix(x);
        }
    }
    void DeleteFix(Node *z){
        for(Node *s; z != root && z->c == Black;){
            int t = (z == z->p->l);
            s = t?z->p->r:z->p->l;
            if(s->c == Red){
                s->c = Black;
                z->p->c = Red;
                t?RotL(z->p):RotR(z->p);
                s = t?z->p->r:z->p->l;
            }
            if(s->l->c == Black && s->r->c == Black){
                s->c = Red;
                z = z->p;
            }
            else{
                if((t?s->r->c:s->l->c) == Black){
                    t ? s->l->c = Black : s->r->c = Black;
                    s->c = Red;
                    t?RotR(s):RotL(s);
                    s = t?z->p->r:z->p->l;
                }
                s->c = z->p->c;
                z->p->c = Black;
                t?s->r->c = Black : s->l->c = Black;
                t?RotL(z->p):RotR(z->p);
                z = root;
            }
        }
        z->c = root->c = Black;
    }
    void Trans(Node *x, Node *y){ (x->p == nullptr) ? root = y : (x == x->p->l) ? x->p->l = y : x->p->r = y; y->p = x->p; }
    void RotL(Node *z){
        Node *y = z->r;
        z->r = y->l;
        if(y->l != nil) y->l->p = z;
        y->p = z->p;

        if(!z->p) root = y;
        else if(z == z->p->l) z->p->l = y;
        else z->p->r = y;
        z->p = y;
        y->l = z;

        z->rn = z->r->ln + z->r->n + z->r->rn;
        y->ln = z->ln + z->n + z->rn;
    }
    void RotR(Node *z){
        Node *x = z->l;
        z->l = x->r;
        if(x->r != nil) x->r->p = z;
        x->p = z->p;

        if(!z->p) root = x;
        else if(z == z->p->l) z->p->l = x;
        else z->p->r = x;
        z->p = x;
        x->r = z;
        
        z->ln = z->l->ln + z->l->n + z->l->rn;
        x->rn = z->ln + z->n + z->rn;
    }
public:
    RBTree() { nil = new Node{0,0,0,0,nullptr,nullptr,nullptr,Black}, root = nil; }
    Node* tree_minimum(Node* t){ for(;t->l != nil;t=t->l); return t; }
    Node* tree_maximum(Node* t){ for(;t->r != nil;t=t->r); return t; }
    void Inorder(Node *t) {
        if (t == nil) return;
        Inorder(t->l);
        std::cout << "(" << t->k << "(" << t->ln << "," << t->n << "," << t->rn << ")," << (t->c == Black ? "B" : "R") << ")";
        Inorder(t->r);
    }
};

int main() {
    ios::sync_with_stdio(0), cin.tie(0);
    RBTree T;
    int n;
    for(cin >> n;n--;){
        int t, x, y; cin >> t;
        if(t==1) cin >> x, cout << T.IsN(x) << endl;
        else cin >> x >> y, y>0 ? T.Insert(x, y) : T.Delete(x,-y);
        //T.Inorder(T.root), cout << endl;
    }
}