Using prefix sum Solution in O(n log(n))


#1
int i,j=0,n=A.size();
vector<long long int> psum(n,0);
for(i=0;i<n;i++)
{
    psum[i]=((i==0)?0:psum[i-1])+A[i]; // prefix sum // sum from 0 till that i;
}

for(i=0;i<n;i++) //for each i as start index, find the end points that satisfy the range
{
    auto l=lower_bound(psum.begin(),psum.end(),(i==0?0:psum[i-1])+B);
    auto r=upper_bound(psum.begin(),psum.end(),(i==0?0:psum[i-1])+C);
    j+=r-l;
    
}
return j;

How to do in O(n)


#2

Hi, I don’t understand your solution, could you explain me what do you do in the second for loop?
In any way thanks.


#3

lower_bound and upper_bound are binary searches. They return the iterators to the respective values. therefore r-l will return the number of elements in that range and subsequently number of continuous arrays with sum between B and C.


#4

Your logic will fail for the following test case
1 2 0 0
0 3
The reason behind it is that lower_bound and upper_bound functions should search from the given value of i and not from the beginning of the prefix sum array.
The answer for this test case should be 10 and not 14 as per your logic.
auto l=lower_bound(psum,begin()+i,psum.end(),(i==0?0:psum[i-1])+B);
auto r=upper_bound(psum.begin()+i,psum.end(),(i==0?0:psum[i-1])+C);


#5

yes you are right.
the code will be
int Solution::numRange(vector &A, int B, int C) {
int i,j,n=A.size();
vectorsum(n);
vector::iterator it;
sum[0]=A[0];
for(i=1;i<n;i++){
sum[i]=sum[i-1]+(long long int)(A[i]);
}
int ans=0;
it=sum.begin();
it++;
ans=ans+upper_bound(sum.begin(),sum.end(),C)-lower_bound(sum.begin(),sum.end(),B);
for(i=1;i<n;i++){
ans=ans+upper_bound(it,sum.end(),C+sum[i-1])-lower_bound(it,sum.end(),B+sum[i-1]);
it++;
}
return ans;
}