Dima and array#2941

Dima and array (Editorial)

The problem requires processing operations of two types:

assigning a value to one element of the array;
computing the sum of the elements on the segment $[f; t]$ .

We need a data structure that performs both operations quickly. Such a structure is a segment tree for sums. Each of its vertices stores the sum of the elements of its segment, and the value in an internal vertex equals the sum of the values of its two children.

The segment tree allows both computing the sum on a segment and changing an array element in $O (lo g n)$ O(\log n).

Example

Let us simulate the queries from the example. Initially, the array is $(1, 2, 3)$ (1, 2, 3).

$? 13$ : the sum of the elements on the segment $[1; 3]$ is $1 + 2 + 3 = 6$ 1 + 2 + 3 = 6;
$= 32$ : the element $a_{3}$ a_3 is assigned the value $2$ , and the array becomes $(1, 2, 2)$ (1, 2, 2);
$? 13$ : the sum of the elements on the segment $[1; 3]$ is $1 + 2 + 2 = 5$ 1 + 2 + 2 = 5.

Let us see how the segment tree changes. The left figure shows the tree for the initial array, and the right figure shows it after the assignment $a_{3} = 2$ a_3 = 2. Each vertex shows its segment and the sum of the elements on it; the changed vertices are highlighted.

The segment $[1; 3]$ of both queries coincides with the segment of the root, so the answer is taken directly from the root: $6$ before the assignment and $5$ after it. During the assignment $a_{3} = 2$ a_3 = 2, only the vertices on the path from the leaf $[3]$ to the root are recomputed: the leaf $[3]$ receives the value $2$ , and the root receives the value $3 + 2 = 5$ 3 + 2 = 5.

Algorithm implementation

Store the segment tree in the vector $S e g T r ee$ SegTree of length $4 \cdot M A X$ 4 \cdot MAX, where $M A X = 5 \cdot 1 0^{5}$ MAX = 5 \cdot 10^5 is the maximum number of array elements. Each vertex of the tree stores the sum of the elements of its segment.

#define MAX 500000
vector<int> SegTree(4*MAX,0);

The function BuildTree builds the segment tree from the array $a$ . It receives the number of the current vertex $v$ and the bounds $l p os$ and $r p os$ of the segment corresponding to this vertex.

void BuildTree(vector<int>&a, int v, int lpos, int rpos)
{

If the segment of vertex $v$ consists of one element ( $l p os = r p os$ lpos = rpos), then $v$ is a leaf of the tree. Write the value $a [l p os]$ into it.

  if (lpos == rpos) SegTree[v] = a[lpos];
  else
  {

Otherwise, find the middle of the segment $mi d = (l p os + r p os) /2$ mid = (lpos + rpos) / 2.

    int mid = (lpos + rpos) / 2;

The segment $[l p os; r p os]$ [lpos; rpos] is split into two subsegments $[l p os; mi d]$ [lpos; mid] and $[mi d + 1; r p os]$ [mid + 1; rpos]. Recursively build the tree for the left child $2 v$ and the right child $2 v + 1$ 2v + 1.

    BuildTree(a, 2*v, lpos, mid);
    BuildTree(a, 2*v+1, mid + 1, rpos);

The value in vertex $v$ equals the sum of the values of its children.

    SegTree[v] = SegTree[2*v] + SegTree[2*v+1];
  }
}

The function Update assigns the value $v a l$ to the element with index $p os$ . Vertex $v$ corresponds to the segment $[l p os; r p os]$ [lpos; rpos].

void Update(int v, int lpos, int rpos, int pos, int val)
{

If vertex $v$ is a leaf ( $l p os = r p os$ lpos = rpos), it corresponds to position $p os$ . Write the value $v a l$ into it.

  if (lpos == rpos) SegTree[v] = val;
  else
  {

Otherwise, determine which half of the segment contains position $p os$ , and recursively descend into the left or right child.

    int mid = (lpos + rpos) / 2;
    if (pos <= mid)
      Update (2*v, lpos, mid, pos, val);
    else
      Update (2*v+1, mid+1, rpos, pos, val);

After returning from the recursion, recompute the sum in vertex $v$ . In this way, all vertices on the path from the leaf to the root are updated.

    SegTree[v] = SegTree[2*v] + SegTree[2*v+1];
  }
}

The function Summa returns the sum of the elements on the segment $[l e f t; r i g h t]$ [left; right]. Vertex $v$ corresponds to the segment $[l p os; r p os]$ [lpos; rpos]. If the query segment is empty ( $l e f t > r i g h t$ left > right), return $0$ — the neutral element for addition.

int Summa(int v, int lpos, int rpos, int left, int right)
{
  if (left > right) return 0;

If the query segment coincides with the segment of the vertex, the answer is stored in $S e g T r ee [v]$ SegTree[v].

  if ((left == lpos) && (right == rpos))
    return SegTree[v];

Otherwise, find the middle of the segment $mi d = (l p os + r p os) /2$ mid = (lpos + rpos) / 2.

  int mid = (lpos + rpos) / 2;

Split the query segment into two parts: $[l e f t; min (r i g h t, mi d)]$ [left; \min(right, mid)] lies in the segment of the left child, and $[max (l e f t, mi d + 1); r i g h t]$ [\max(left, mid + 1); right] lies in the segment of the right child. One of the parts may be empty. Return the sum of the results for both children.

  return Summa(2*v, lpos, mid, left, min(right,mid)) +
         Summa(2*v+1, mid+1, rpos, max(left,mid+1), right);
}

The main part of the program. Read the size of the array $n$ , the number of queries $q$ , and the array elements into the vector $v$ , starting from index $1$ . The call scanf("\n") skips the line break after the array, so that the next character read is the operation sign.

scanf("%d %d",&n,&q);
v.resize(n+1);
for(i = 1; i <= n; i++)
  scanf("%d",&v[i]);
scanf("\n");

Build the segment tree from the elements $v [1 \dots n]$ v[1 \ldots n].

BuildTree(v,1,1,n);

Process the $q$ queries one by one. First, read the operation character $c$ ; the space in the format "%c " skips the space after it. If $c$ = '=', assign the value $d$ to the element with index $i$ . Otherwise, compute and print the sum on the segment $[f; t]$ . The \n character in the format "%d %d\n" skips the line break, so the next read into $c$ gets the sign of the next operation.

for(j = 0; j < q; j++)
{
  scanf("%c ",&c);
  if (c == '=')
  {
    scanf("%d %d\n",&i,&d);
    Update(1,1,n,i,d);
  } else
  {
    scanf("%d %d\n",&f,&t);
    printf("%d\n",Summa(1,1,n,f,t));
  }
}

C++