algorithmdata-structuressegment-tree

How can I solve the problem on the segment tree

I have an array a[0..n - 1]. I have 2 types operations:

  1. sum l r: sum = a[l] * 1 + a[l+1] * 2 + ... + a[r] * (r - l)
  2. update index value: a[index] = value

How can I modify standart sum segment tree for this task?

Solution

  • Store a second segment tree for sums of the array b, where b[i] = (i+1) * a[i]. Then, setting a[index] = value should also update the second segment tree with b[index] = value * (index + 1).

    To calculate the sum you want, given the ability to query the normal subarray sums of a and b:

    special_sum(l, r) = a[l] * 1 + a[l+1] * 2 + ... + a[r] * (r - l)
can be rewritten as:

(l+1) * a[l] + (l+2) * a[l+1] + ... + (r) * a[r]
- l * (a[l] + a[l+1] + ... + a[r])

which becomes

b[l] + b[l+1] + ... + b[r]
- l * (a[l] + a[l+1] + ... + a[r])