Suma en un segmento (Editorial)

El problema requiere procesar consultas de dos tipos: asignar un valor a un elemento del arreglo y calcular la suma de los elementos en el segmento $[l; r]$ [l; r]. Si almacenamos solo el arreglo en sí, una asignación toma $O (1)$ , pero cada consulta de suma requiere $O (n)$ operaciones. Si almacenamos sumas de prefijos, una suma se calcula en $O (1)$ , pero cada asignación requiere recalcular $O (n)$ sumas de prefijos. Con una gran cantidad de consultas, ambos enfoques son demasiado lentos.

Usamos un árbol de segmentos para las sumas. Cada uno de sus vértices almacena la suma de los elementos de su segmento, y el valor en un vértice interno es igual a la suma de los valores de sus dos hijos. El árbol de segmentos realiza ambas operaciones en $O (lo g n)$ O(\log n). Al calcular una suma, una parte vacía de un segmento corresponde a $0$ — el elemento neutro para la suma.

Inicialmente todos los elementos del arreglo son cero, por lo que las sumas en todos los vértices del árbol también son cero. Basta con llenar el arreglo del árbol con ceros; no es necesario construir el árbol por separado.

Los valores de los elementos alcanzan $1 0^{9}$ 10^9, por lo que la suma en un segmento puede alcanzar $1 0^{5} \cdot 1 0^{9} = 1 0^{14}$ 10^5 \cdot 10^9 = 10^{14} y no cabe en el tipo int. Los valores en el árbol se almacenan en el tipo long long.

Consideremos una solución usando un árbol de Fenwick. Admite dos operaciones, cada una realizada en $O (lo g n)$ O(\log n):

sumar un número $d e l t a$ al elemento $a_{i}$ a_i;
calcular la suma de prefijo $s u m (i) = a_{0} + a_{1} + \dots + a_{i}$ sum(i) = a_0 + a_1 + \ldots + a_i.

Inicialmente todos los elementos del arreglo son cero, por lo que un árbol de Fenwick lleno de ceros ya corresponde al arreglo inicial. Los índices en las consultas van de $1$ a $n$ ; el elemento $a_{0}$ a_0 no se usa y siempre es cero.

Asignación. Un árbol de Fenwick solo puede sumar un número a un elemento, pero no puede asignarle un nuevo valor. Sin embargo, el valor actual de un elemento se puede calcular mediante sumas de prefijos:

a_{i} = s u m (i) - s u m (i - 1)

a_i = sum(i) - sum(i - 1)

Por lo tanto, la asignación $a_{i} = x$ a_i = x se reemplaza sumando a $a_{i}$ a_i la diferencia entre el nuevo y el antiguo valor:

d e l t a = x - (s u m (i) - s u m (i - 1)) = x - s u m (i) + s u m (i - 1)

delta = x - (sum(i) - sum(i - 1)) = x - sum(i) + sum(i - 1)

Suma en un segmento. La suma de los elementos en el segmento $[l; r]$ [l; r] es igual a la diferencia de dos sumas de prefijos:

a_{l} + a_{l + 1} + \dots + a_{r} = s u m (r) - s u m (l - 1)

a_l + a_{l+1} + \ldots + a_r = sum(r) - sum(l - 1)

Los valores de los elementos alcanzan $1 0^{9}$ 10^9, por lo que las sumas pueden alcanzar $1 0^{14}$ 10^{14} y se almacenan en el tipo long long.

Ejemplo

Inicialmente el arreglo consiste en ceros. Después de las consultas $A 22$ , $A 31$ y $A 42$ , el arreglo se convierte en $(0, 2, 1, 2, 0)$ (0, 2, 1, 2, 0). El árbol de segmentos para él se muestra a continuación. Cada vértice muestra su segmento y la suma de los elementos en él.

Cuando se ejecuta la consulta $A 31$ , solo se recalculan los vértices en el camino desde la hoja $[3]$ hasta la raíz: la hoja $[3]$ recibe el valor $1$ , el vértice $[1; 3]$ — el valor $2 + 1 = 3$ 2 + 1 = 3, y la raíz $[1; 5]$ — el valor $3 + 2 = 5$ 3 + 2 = 5.

Las consultas $Q 11$ , $Q 22$ , $Q 33$ , $Q 44$ y $Q 55$ devuelven los valores de las hojas: $0, 2, 1, 2, 0$ 0, 2, 1, 2, 0. El segmento de la consulta $Q 15$ coincide con el segmento de la raíz, por lo que la respuesta se toma directamente de la raíz y es igual a $5$ (el vértice resaltado).

Simulemos las consultas del ejemplo usando el árbol de Fenwick, $n = 5$ n = 5. Inicialmente el arreglo es $(0, 0, 0, 0, 0)$ (0, 0, 0, 0, 0). Para cada consulta de asignación, la tabla muestra cómo se calcula el valor antiguo del elemento y qué número se suma; para cada consulta de suma, muestra la diferencia de sumas de prefijos.

consulta	cálculo	resultado
$A 22$	valor antiguo $s u m (2) - s u m (1) = 0 - 0 = 0$ sum(2) - sum(1) = 0 - 0 = 0, se suma $2 - 0 = 2$ 2 - 0 = 2	$(0, 2, 0, 0, 0)$ (0, 2, 0, 0, 0)
$A 31$	valor antiguo $s u m (3) - s u m (2) = 2 - 2 = 0$ sum(3) - sum(2) = 2 - 2 = 0, se suma $1 - 0 = 1$ 1 - 0 = 1	$(0, 2, 1, 0, 0)$ (0, 2, 1, 0, 0)
$A 42$	valor antiguo $s u m (4) - s u m (3) = 3 - 3 = 0$ sum(4) - sum(3) = 3 - 3 = 0, se suma $2 - 0 = 2$ 2 - 0 = 2	$(0, 2, 1, 2, 0)$ (0, 2, 1, 2, 0)
$Q 11$	$s u m (1) - s u m (0) = 0 - 0$ sum(1) - sum(0) = 0 - 0	$0$
$Q 22$	$s u m (2) - s u m (1) = 2 - 0$ sum(2) - sum(1) = 2 - 0	$2$
$Q 33$	$s u m (3) - s u m (2) = 3 - 2$ sum(3) - sum(2) = 3 - 2	$1$
$Q 44$	$s u m (4) - s u m (3) = 5 - 3$ sum(4) - sum(3) = 5 - 3	$2$
$Q 55$	$s u m (5) - s u m (4) = 5 - 5$ sum(5) - sum(4) = 5 - 5	$0$
$Q 15$	$s u m (5) - s u m (0) = 5 - 0$ sum(5) - sum(0) = 5 - 0	$5$

Por ejemplo, antes de la consulta $Q 44$ el arreglo es $(0, 2, 1, 2, 0)$ (0, 2, 1, 2, 0), por lo que $s u m (4) = 0 + 2 + 1 + 2 = 5$ sum(4) = 0 + 2 + 1 + 2 = 5, $s u m (3) = 0 + 2 + 1 = 3$ sum(3) = 0 + 2 + 1 = 3, y la respuesta es $5 - 3 = 2$ 5 - 3 = 2.

En el ejemplo, todas las asignaciones se aplican a elementos nulos. Veamos qué pasaría con una asignación repetida, por ejemplo $A 35$ al final. El valor antiguo es $s u m (3) - s u m (2) = 3 - 2 = 1$ sum(3) - sum(2) = 3 - 2 = 1, por lo que se suma $5 - 1 = 4$ 5 - 1 = 4 a $a_{3}$ a_3, y el elemento se convierte en $5$ .

Implementación del algoritmo

El árbol de segmentos se almacena en el arreglo $S e g T r ee$ SegTree de longitud $4 \cdot M A X$ 4 \cdot MAX, donde $M A X$ es el número máximo de elementos del arreglo. Una suma en un segmento puede exceder los límites del tipo int, por lo que los valores en el árbol tienen el tipo long long.

#define MAX 100010
long long SegTree[4 * MAX];

La función Summa devuelve la suma de los elementos en el segmento $[l e f t; r i g h t]$ [left; right]. El vértice $v$ corresponde al segmento $[l p os; r p os]$ [lpos; rpos].

long long Summa(int v, int lpos, int rpos, int left, int right)
{

Si el segmento de la consulta está vacío ( $l e f t > r i g h t$ left > right), devuelve $0$ — el elemento neutro para la suma.

  if (left > right) return 0;

Si el segmento de la consulta coincide con el segmento $[l p os; r p os]$ [lpos; rpos] del vértice $v$ , devuelve la suma almacenada en él.

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

De lo contrario, encuentra el punto medio del segmento $mi d = (l p os + r p os) /2$ mid = (lpos + rpos) / 2.

  int mid = (lpos + rpos) / 2;

Divide el segmento de la consulta en dos partes: $[l e f t; min (r i g h t, mi d)]$ [left; \min(right, mid)] se encuentra en el segmento del hijo izquierdo, y $[max (l e f t, mi d + 1); r i g h t]$ [\max(left, mid + 1); right] se encuentra en el segmento del hijo derecho. Una de las partes puede estar vacía. Calcula recursivamente las sumas en ambas partes y devuelve su suma.

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

La función update asigna el valor $v a l$ al elemento con índice $p os$ . El vértice $v$ corresponde al segmento $[l p os; r p os]$ [lpos; rpos].

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

Si el segmento del vértice $v$ consiste en un solo elemento ( $l p os = r p os$ lpos = rpos), entonces $v$ es la hoja correspondiente a la posición $p os$ . Escribe el valor $v a l$ en ella.

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

De lo contrario, determina qué mitad del segmento contiene la posición $p os$ , y desciende recursivamente al hijo izquierdo o derecho.

    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);

Después de regresar de la recursión, recalcula la suma en el vértice $v$ . De esta manera, se actualizan todos los vértices en el camino desde la hoja hasta la raíz.

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

La parte principal del programa. Inicialmente todos los elementos del arreglo son cero, por lo que todas las sumas en el árbol de segmentos también son cero. Llena el arreglo $S e g T r ee$ SegTree con ceros; no es necesario construir el árbol por separado.

memset(SegTree, 0, sizeof(SegTree));

Lee el tamaño del arreglo $n$ y el número de consultas $k$ . El carácter \n en el formato omite el salto de línea, para que el siguiente carácter leído sea la letra de la consulta.

scanf("%d %d\n", &n, &k);

Procesa las $k$ consultas una por una. Primero, lee la letra de la consulta $co mman d$ . Si es 'A', asigna el valor $x$ al elemento con índice $i$ . De lo contrario, imprime la suma en el segmento $[l; r]$ [l; r]. El carácter \n en el formato "%lld %lld\n" omite el salto de línea después de la consulta.

for (j = 0; j < k; j++)
{
  scanf("%c", &command);
  if (command == 'A')
  {
    scanf("%lld %lld\n", &i, &x);
    update(1, 1, n, i, x);
  }
  else
  {
    scanf("%lld %lld\n", &l, &r);
    printf("%lld\n", Summa(1, 1, n, l, r));
  }
}

C++

Implementación del algoritmo — árbol de Fenwick

El arreglo $F e n w i c k$ Fenwick de tamaño $M A X$ almacena el árbol de Fenwick. Un arreglo global se llena inicialmente con ceros, lo que corresponde al arreglo inicial nulo.

#define MAX 100010
long long Fenwick[MAX];

La función Summma0_i devuelve la suma de prefijo $s u m (i) = a_{0} + a_{1} + \dots + a_{i}$ sum(i) = a_0 + a_1 + \ldots + a_i.

// Fenwick[0] + Fenwick[1] + ... + Fenwick[i]
long long Summma0_i(long long i)
{
  long long result = 0;
  for (; i >= 0; i = (i & (i + 1)) - 1)
    result += Fenwick[i];
  return result;
}

C++

La función IncElement suma $d e l t a$ al elemento $a_{i}$ a_i. Los índices de los elementos van de $1$ a $n$ , por lo que el ciclo se ejecuta mientras $i \leq n$ i \le n.

// Fenwick[i] = Fenwick[i] + delta
void IncElement(long long i, long long delta)
{
  for (; i <= n; i = (i | (i+1)))
    Fenwick[i] += delta;
}

La parte principal del programa. Lee el tamaño del arreglo $n$ y el número de consultas $k$ . El carácter \n en el formato omite el salto de línea, para que el siguiente carácter leído sea la letra de la consulta.

scanf("%d %d\n",&n,&k);

Procesa las $k$ consultas una por una. Primero, lee la letra de la consulta $co mman d$ .

for(j = 0; j < k; j++)
{
  scanf("%c",&command);

Consulta $A i x$ . El valor actual del elemento $a_{i}$ a_i es igual a $s u m (i) - s u m (i - 1)$ sum(i) - sum(i - 1). Para asignarle el valor $x$ , reemplaza $x$ por la diferencia $x - s u m (i) + s u m (i - 1)$ x - sum(i) + sum(i - 1) y súmala a $a_{i}$ a_i usando la función IncElement.

  if (command == 'A')
  {
    scanf("%lld %lld\n",&i,&x);
    x = x - Summma0_i(i) + Summma0_i(i-1);
    IncElement(i,x);

Consulta $Q l r$ . La suma en el segmento $[l; r]$ [l; r] es igual a $s u m (r) - s u m (l - 1)$ sum(r) - sum(l - 1).

  } else
  {
    scanf("%lld %lld\n",&l,&r);
    printf("%lld\n",Summma0_i(r) - Summma0_i(l-1));
  }
}