weakestOsuPlayer_244's blog

Searching for an array prefix with a given amount. (using segment trees)
By weakestOsuPlayer_244, history, 5 years ago, In English

While learning about segment tree I came across this problem:-

The task is as follows: for a given value x we have to quickly find smallest index i such that the sum of the first i elements of the array a[] is greater or equal to x (assuming that the array a[] only contains non-negative values).

Link-https://cp-algorithms.com/data_structures/segment_tree.html

The solution mentions storing prefix sums in the segment trees.

I cannot figure out how to build up a segment tree of prefix sums. Also how is it different from the range sums?

  • Vote: I like it
  • +3
  • Vote: I do not like it

Comments Comments (7)
| Write comment?
»
5 years ago, # |
← Rev. 2   Vote: I like it +4 Vote: I do not like it

No, we don't need to store the prefix sum. We just need to build the simple sum segment tree, it only differs in the way we query the segment tree. Let the required sum be K, idx represent the index of the segment tree and l and r be the array range it covers, then code will look like,

int query(int idx, int l, int r, int K){
    if( l == r )
        return l;
    int md = (l+r)/2;
    if( seg[2*idx] >= K ) // left child contains the required index
        return query(2*idx, l, md, K);
    else
        return query(2*idx+1, md+1, r, K-seg[2*idx]); // right segment don't have sum of left segment
}

/*********************************
if( seg[1] < K ){
    // answer doesn't exist
}else{
    ans = query(1, 0, n-1, K);
}
********************************/
→ Reply
  • »
    »
    5 years ago, # ^ |
      Vote: I like it 0 Vote: I do not like it

    return query(2*idx+1, md+1, r, K-seg[2*idx]);

    This is how I understood your code. Supposing the left child doesn't have the sum>=K so we go to right child and look for the K-left child sum. Now if the right child's sum is greater than K-left child sum the l of the right child will be the answer. (since the sum of left child +this right child will end up greater than K)

    This is how it is?

    → Reply
    • »
      »
      »
      5 years ago, # ^ |
      ← Rev. 2   Vote: I like it 0 Vote: I do not like it

      Yes, since we know answer do exist (see comments in the code), then if answer is not in the left child then answer will be in the segment covered by right child and hence we call the same function on the right child with sum of the left child removed

      → Reply
»
5 years ago, # |
  Vote: I like it 0 Vote: I do not like it

Do you need to update the elements of the array?

If no, use binary search to find the first position.

My code (C++)
Sample tests

Otherwise, my solution would be to use a segment tree with range sums.

My code(C++)
Sample tests
→ Reply
  • »
    »
    5 years ago, # ^ |
      Vote: I like it +3 Vote: I do not like it

    Thanks for the detailed implementation. Yes, I was looking for a O(log n) implementation using segment tree.

    → Reply
»
3 years ago, # |
  Vote: I like it 0 Vote: I do not like it 
template<typename T>int find_prefix_seg(vt & t,int v,int tl,int tr,int x){
	if(x>t[v]) return -1;
	if(tl==tr){
		return tl;
	}
	int tm=MID(tl,tr);
	if(t[2*v+1]>=x) return find_prefix_seg(t,v*2+1,tl,tm,x);
	return find_prefix_seg(t,v*2+2,tm+1,tr,x-t[v*2+1]);
}
→ Reply