KQuery SPOJ problem, using merge sort segment tree

Question

Problem : KQUERY

Summary: For each query consisting of \$(i, j, k)\$, return the number of elements greater than \$k\$ in the subsequence \$a_i, a_{i+1}, …, a_j\$.

My code has complexity of \$O(n\log^2{n})\$. I am only interested in segment tree based solution of Kquery. Please let me know why am I getting TLE using merge sort segment tree. Is it possible to do some optimization to get my code accepted or is this approach totally wrong. Here I have used merge sort segment tree which has sorted vector as its node and to query particular range it creates the sorted vector of elements for the given query.

#include <bits/stdc++.h>
using namespace std;

const int N=30005;
int arr[N];
vector<int> seg[4*N];

void build(int low,int high,int node)
{
    if(low>high)
    return;
    if(low == high)
    {
        seg[node].push_back(arr[low]);
        return;
    }
    int mid=low+high>>1;
    build(low,mid,2*node+1);
    build(mid+1,high,2*node+2);
    merge(seg[2*node+1].begin(),seg[2*node+1].end(),seg[2*node+2].begin(),seg[2*node+2].end(),back_inserter(seg[node]));
}

vector<int> query(int low,int high,int lq,int hq,int k,int node)
{
    vector<int> ans;
    if(low>high || low>hq || high<lq)
    return ans;
    if(lq<=low && high<=hq)
    {
        return seg[node];
    }
    vector<int> left,right;
    int mid=low+high>>1;
    left = query(low,mid,lq,hq,k,2*node+1);
    right = query(mid+1,high,lq,hq,k,2*node+2);
    merge(left.begin(),left.end(),right.begin(),right.end(),back_inserter(ans));
    return ans;
}
int main(){
    int n;
    scanf("%d",&n);
    register int i;
    for(i=0;i<n;i++)
    scanf("%d",&arr[i]);
    register int q;
    scanf("%d",&q);
    int x,y,k;
    build(0,n-1,0);
    while(q--)
    {
        scanf("%d %d %d",&x,&y,&k);
        vector<int> ans=query(0,n-1,x-1,y-1,k,0);
        int point=upper_bound(ans.begin(),ans.end(),k)-ans.begin();
        printf("%d\n",ans.size()-point);
    }
    return 0;
}

Answer 1

Your biggest issue here is that you are building a new list for every single query. Each of these lists can have size up to \$N\$ (since the query could be over the entire universe), so building this list takes time \$O(N)\$ in the worst case (this is true even though it is built from logarithmically many smaller lists). This makes the total runtime of your algorithm \$O(N\log^2 N + QN)\$, and \$O(QN)\$ is way too large to run in time for the parameters of this problem.

Of course, the solution to this issue is that you don't really need to build the list for each query; you just need to count how many elements are larger than k. Since the list for each query is just the union of logarithmically many sorted lists in your segment tree, just count this quantity for each of those lists and sum this up. This reduces your time complexity to \$O(N\log^2 N + Q\log^2 N)\$ (there are \$O(\log N)\$ lists, and it takes \$O(\log N)\$ time to binary search each one), which should run in time.

If your code still times out after making this change, I would recommend rewriting the code for build and query iteratively instead of recursively; this can sometimes speed things up a bit (however we're no longer in the realm of asymptotic improvements).

Some other code things:

• I might be wrong, but I think it's generally accepted that the register keyword in C++ doesn't really do much anymore; compilers are way better than people at figuring out which variables should be stored in registers, so you're not going to get any speedup by explicitly declaring this.
• Your query method accepts k as a parameter, but you don't actually use k anywhere (short of passing it along in recursive calls). That said, given the above answer, you probably should be using k in the query method.
• I guess it's not super necessary for short and quick solutions to contest problems like this, but I think your code would benefit a lot from a bit more whitespace and consistent indentation (I think it would make it a lot easier to spot bugs as well). One of the things that bugs me the most about your code is how you write for(i=0;i<n;i++) and scanf("%d", &arr[n]) on two consecutive lines at the same indentation; it makes it very unclear that the second line is part of the for loop.