C++ Treap O(NlogN)


  • 0
    J
    template<class T>
    struct TreapSet {
        static unsigned state;
        TreapSet *l, *r;
        unsigned priority;
        T key;
        int size;
        TreapSet(T _key) : l(nullptr), r(nullptr), priority(state), key(_key), size(1) {
            state = state * 1297 + 1301;
        }
        void update() {
            size = 1;
            if (l) size += l->size;
            if (r) size += r->size;
        }
        ~TreapSet() {
            if (l) delete l;
            if (r) delete r;
        }
    };
    
    template<>
    unsigned TreapSet<long long>::state = 0;
    
    template<class T>
    TreapSet<T> *merge(TreapSet<T> *a, TreapSet<T> *b) {
        if (!a || !b) return a ? a : b;
        if (a->priority > b->priority) {
            a->r = merge(a->r, b);
            a->update();
            return a;
        } else {
            b->l = merge(a, b->l);
            b->update();
            return b;
        }
    }
    
    template<class T>
    void split(TreapSet<T> *t, T k, TreapSet<T> *&a, TreapSet<T> *&b) {
        if (!t) a = b = nullptr;
        else if (t->key <= k) {
            a = t;
            split(t->r, k, a->r, b);
            a->update();
        } else {
            b = t;
            split(t->l, k, a, b->l);
            b->update();
        }
    }
    
    template<class T>
    void insert(TreapSet<T> *&t, T key) {
        TreapSet<T> *a, *b;
        split(t, key, a, b);
        t = merge(a, merge(new TreapSet<T>(key), b));
    }
    
    class Solution {
    public:
        int countRangeSum(vector<int>& nums, int lower, int upper) {
            TreapSet<long long> *t = new TreapSet<long long>(0);
            long long prefix = 0, ans = 0;
            for (int num : nums) {
                prefix += num;
                TreapSet<long long> *a, *bc, *b, *c;
                split(t, prefix - upper - 1, a, bc);
                split(bc, prefix - lower, b, c);
                if (b) ans += b->size;
                t = merge(a, merge(b, c));
                insert(t, prefix);
            }
            delete t;
            return ans;
        }
    };
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.