Allin1

I’m trying to solve ALLIN1 problem using Treap (as it is given in cp-algorithms in that section). I’m getting TLE and as I read the comments, it’s not only me. I cannot figure out a better designed solution using Treaps…

Can anyone help, is there a solution with Treaps?

Here is mine…

//  https://www.spoj.com/problems/ALLIN1/

// Time complexity: all queries are O(log(N))
// Space complexity: O(N)

// Treap supports duplicate keys (multiset)

#include <cstdio>
#include <cstdlib>
#include <ctime>

using namespace std;

typedef struct Node *pnode;

struct Node {
    int key, prior, cnt, sub; // cnt = count of duplicates for keys in current Node, sub = subtree count 
    pnode l, r;

    Node(int key) : key(key), prior(rand()), cnt(1), sub(1), l(nullptr), r(nullptr) {}
};

struct Treap {
private:
    pnode root;

    void clear(pnode t) {
        if (t) {
            clear(t->l);
            clear(t->r);
            delete t;
        }
    }

    int sub(pnode t) {
        return t ? t->sub : 0;
    }

    void upd(pnode t) {
        if (t) {
            t->sub = t->cnt + sub(t->l) + sub(t->r);
        }
    }

    void split(pnode t, int key, pnode &l, pnode &r) {
        if (!t) {
            l = r = nullptr;
        } else if (key < t->key) {
            split(t->l, key, l, t->l), r = t;
        } else {
            split(t->r, key, t->r, r), l = t;
        }
        upd(t);
    }

    void merge(pnode &t, pnode l, pnode r) {
        if (!l || !r) {
            t = l ? l : r;
        } else if (l->prior > r->prior) {
            merge(l->r, l->r, r), t = l;
        } else {
            merge(r->l, l, r->l), t = r;
        }
        upd(t);
    }

    pnode find(pnode t, int key) {
        if (!t) {
            return nullptr;
        }
        if (t->key == key) {
            return t;
        }
        if (key < t->key) {
            return find(t->l, key);
        } else {
            return find(t->r, key);
        }
    }

    void updPath(pnode t, int key) {
        if (!t) {
            return;
        }
        if (t->key == key) {}
        else if (key < t->key) {
            updPath(t->l, key);
        } else {
            updPath(t->r, key);
        }
        upd(t);
    }

    void insert(pnode &t, pnode it) {
        if (!t) {
            t = it;
        } else if (t->prior > it->prior) {
            insert(t->key <= it->key ? t->r : t->l, it);
        } else {
            split(t, it->key, it->l, it->r), t = it;
        }
        upd(t);
    }

    void erase(pnode &t, int key) {
        if (!t) {
            return;
        }
        if (t->key == key) {
            pnode toDel = t;
            merge(t, t->l, t->r);
            delete toDel;
            upd(t);
        } else {
            erase(t->key < key ? t->r : t->l, key);
            upd(t);
        }
    }

    int rank(pnode t, int key) {
        int rank = 1;
        while (t) {
            if (key < t->key) {
                t = t->l;
            } else if (key > t->key) {
                rank += sub(t->l) + t->cnt;
                t = t->r;
            } else {
                return rank + sub(t->l);
            }
        }
        return -1;
    }

    int kthSmallest(pnode t, int k) {
        if (sub(t) < k) {
            return -1;
        }
        while (t) {
            int lst = sub(t->l);

            if (lst >= k) {
                t = t->l;
            } else if (lst + t->cnt >= k) {
                return t->key;
            } else {
                k -= lst + t->cnt;
                t = t->r;
            }
        }
        return -1;  // we will never reach here
    }

public:
    Treap() : root(nullptr) {
        srand(time(nullptr));
    }

    ~Treap() {
        clear(this->root);
    }

    void smartInsert(int key) { // if it exist - increment and update Path, else insert new element
        pnode found = find(this->root, key);
        if (!found) {
            insert(this->root, new Node(key));
        } else {
            ++found->cnt;
            updPath(this->root, key);
        }
    }

    void smartErase(int key) {
        pnode found = find(this->root, key);
        if (found) {
            int freq = --found->cnt;
            if (!freq) {
                erase(this->root, key);
            }
            updPath(this->root, key);
        }
    }

    int rank(int key) {  // rank = first index of key if all keys were sorted
        return rank(this->root, key);
    }

    int kthSmallest(int k) {
        return kthSmallest(this->root, k);
    }
};


int main() {
    Treap t;

    int type, x;
    while (scanf("%d", &type) && type != -1) {
        scanf("%d", &x);
        switch (type) {
            case 1: {
                t.smartInsert(x);
                break;
            }

            case 2: {
                t.smartErase(x);
                break;
            }

            case 3: {
                printf("%d\n", t.rank(x));
                break;
            }

            case 4: {
                printf("%d\n", t.kthSmallest(x));
                break;
            }

            default: {
                printf("ivalid query type\n");
            }
        }
    }
    return 0;
}