basic TREAP, C++ (gcc)

basic TREAP

#include <bits/stdc++.h>

using namespace std;

typedef long long int ll;
#define mp make_pair
#define pb push_back
#define pii pair<int, int>

struct Treap{
    struct node{
        node* lft, *rgt;
        ll prior, val;
        int cnt;
        
        node(){}
        node(ll v, ll p, node* l = NULL, node* r = NULL) : val(v), prior(p), lft(l), rgt(r){}
    };
    
    typedef node* pnode;
    
    pnode root;
    
    Treap() : root(NULL){}
    
    int get_cnt(pnode p){
        return (p == NULL)? 0 : p -> cnt;
    }
    
    void upd_cnt(pnode p){
        if(p) p -> cnt = get_cnt(p -> lft) + get_cnt(p -> rgt) + 1;
    }
    
    void split(pnode root, pnode& l, pnode& r, int ind){
        if(!root) return void(l = r = NULL);
        int cur = 1 + get_cnt(root -> lft);
        
        if(cur > ind) split(root -> lft, l, root -> lft, ind), r = root;
        else split(root -> rgt, root -> rgt, r, ind - cur), l = root;
        
        //upd_cnt(l), upd_cnt(r);
        upd_cnt(root);
    }
    
    void merge(pnode& root, pnode l, pnode r){
        if(!l || !r) return void(root = l? l : r);
        if(l -> prior > r -> prior) merge(l -> rgt, l -> rgt, r), root = l;
        else merge(r -> lft, l, r -> lft), root = r;
        
        upd_cnt(root);
    }
    
    pnode insert(pnode root, pnode cur, int pos){
        if(!root) return cur;
        int at = get_cnt(root -> lft) + 1;
        
        if(root -> prior < cur -> prior || at == pos){
            pnode l, r;
            split(root, l, r, pos-1);
            cur -> lft = l, cur -> rgt = r;
            upd_cnt(cur);
            return cur;
        }
        (at > pos)? root -> lft = insert(root -> lft, cur, pos) : root -> rgt = insert(root -> rgt, cur, pos - at);
        upd_cnt(root -> lft), upd_cnt(root -> rgt), upd_cnt(root);
        return root;
    }
    
    void insert(int pos, int val){
        ll p = rand();
        pnode cur = new node(val, p);
        cur -> cnt = 1;
        
        if(!root) return void(root = cur);
        pnode t1, t2;
        split(root, t1, t2, pos-1);
        merge(t1, t1, cur);
        merge(root, t1, t2);
        
        //root = insert(root, cur, pos);        // this thing is less efficient than the original merge and split thing!
    }
    
    pnode erase(pnode root, int pos){
        if(!root) return NULL;
        int cur = get_cnt(root -> lft) + 1;
        if(pos == cur){
            pnode ret;
            merge(ret, root -> lft, root -> rgt);
            upd_cnt(ret);
            return ret;
        }
        if(pos < cur) root -> lft = erase(root -> lft, pos);
        else root -> rgt = erase(root -> rgt, pos-cur);
        upd_cnt(root -> lft), upd_cnt(root -> rgt), upd_cnt(root);
        return root;
    }
    
    void erase(int pos){
        root = erase(root, pos);
    }
    
    int at(pnode root, int pos){
        if(!root) return -1;
        int cur = get_cnt(root -> lft) + 1;
        if(pos == cur) return root -> val;
        return (pos < cur)? at(root -> lft, pos) : at(root -> rgt, pos-cur);
    }
    
    int at(int pos){
        return at(root, pos);
    }
    
    void preorder(pnode root){
        if(!root) return;
        cout << root -> val << " ";
        preorder(root -> lft);
        preorder(root -> rgt);
    }
    
    void print(pnode root){
        if(!root) return;
        print(root -> lft);
        cout << root -> val << " ";
        print(root -> rgt);
    }
    
    void print(){
        cout << "inorder : ";
        print(root);
        cout << endl << "preorder : ";
        preorder(root);
        cout << endl << endl;
    }
};

int main(){
    ios_base :: sync_with_stdio(0);
    cin.tie(0);
    
    int n, q;
    scanf("%d %d",&n ,&q);
    Treap tp;
    for(int j = 1; j <= n; j++){
        int x;
        scanf("%d",&x);
        tp.insert(j, x);
        //tp.print();
    }
    while(q--){
        int typ;
        scanf("%d",&typ);
        if(typ == 1){
            int k, x;
            scanf("%d %d",&k ,&x);
            tp.insert(k, x);
            //tp.print();
        }
        else if(typ == 2){
            int k;
            scanf("%d",&k);
            tp.erase(k);
            //tp.print();
        }
        else{
            int k;
            scanf("%d",&k);
            printf("%d\n",tp.at(k));
        }
    }
    return 0;
}


run \| edit \| history \| help	0

Please log in to post a comment.