Triple sum

  • + 0 comments

    Completely unnecessary Treap implementation in C++

    This can be solved in 2 lines with python, but why do it when we can implement a binary search tree in C++ ourselves :)

    Python:

    a,c,b = sorted(set(a)), sorted(set(c)),set(b)
        return sum([bisect(a,i)*bisect(c,i) for i in b])
    

    C++:

    struct node{
        int v,p;
        uint64_t c;
        node* left, *right;
        node(int _v, int _p){c=1;v=_v;p=_p;left=right=nullptr;};
        inline static uint64_t getCount(node* x) {
            if(!x) return 0;
            return 1+(x->left?x->left->c:0)+(x->right?x->right->c:0); 
        }
        void update() {c=1+getCount(left)+getCount(right);}
        static void split(node* p, node*& l, node*& r, int val){
            if(!p) l=r=p;
            else if(p->v < val) split(p->right, p->right,r,val), l=p;
            else split(p->left, l, p->left, val), r=p;
            if(p) p->update();
        }
        static void insert(node*& p, node* x){
            if(!p) p=x;
            else if(x->p > p->p) split(p, x->left, x->right, x->v), p=x;
            else (p->v > x->v ? insert(p->left,x) : insert(p->right,x) );
            if(p) p->update();
        }
        static uint64_t lessThan(node* p, int val){
            if(!p) return 0;
            else if (p->v > val) return lessThan(p->left,val);
            else return 1+lessThan(p->right,val)+getCount(p->left);
        }
    };
    
    uint64_t triplets(vector<int> a, vector<int> b, vector<int> c) {
        sort(a.begin(),a.end());
        a.resize(distance(a.begin(),unique (a.begin(), a.end())));
        unordered_set<int> sa, sb, sc;
        node *ax=nullptr, *cx=nullptr;
        uint64_t ans = 0;
        for(auto x: a) node::insert(ax,new node(x,rand()));
        for(auto x: c) if(sc.insert(x).second) 
            node::insert(cx,new node(x,rand()));
        for(auto x: b) if(sb.insert(x).second) 
            ans+=node::lessThan(ax,x)*node::lessThan(cx,x);
        return ans;
    }