Lenja propagacija za sumu u segmentu

Nad nizom brojeva podržati dve operacije pomoću segmentnog stabla sa lenjom propagacijom:

  1. uvećanje svih elemenata segmenta [a, b] za istu vrednost,
  2. računanje sume elemenata u segmentu [a, b].

Opis ulaza

Sa standardnog ulaza se unosi ceo broj n, a zatim i n celih brojeva. Nakon toga, unosi se broj q, a zatim i q upita. Svaki od upita je sledećeg oblika:

Opis izlaza

Ispisati odgovarajuću sumu za svaki od upita drugog tipa.

Primer

Ulaz

8 1 2 3 4 5 6 7 8 5 u 2 5 4 u 1 3 7 q 2 4 u 5 5 3 q 4 7

Izlaz

38 37

Rešenje

Opis glavnog rešenja

Klasično segmentno stablo je prošireno dodatnim nizom lazy_values koji čuva odložena ažuriranja za čvorove čija podstabla još nisu eksplicitno ažurirana. Kada se nad nekim segmentom zatraži povećanje svih elemenata za diff, čvor koji ceo segment pokriva može odmah da ažurira svoju sumu, dok se samo informacija o odloženom povećanju prosledi potomcima.

Na taj način se izbegava obilazak svih listova pri svakom ažuriranju opsega. I ažuriranje opsega i upit sume rade u O(log n), uz linearnu memoriju za stablo i lazy niz.

#include <iostream>
#include <vector>
#include <cmath>

// Rekurzivna funkcija koja kreira segmentno stablo segment_tree od niza brojeva array, k predstavlja poziciju cvora unutar niza (njegov indeks), dok x i y oznacavaju krajeve intervala
// koje cvor obuhvata, na pocetku je taj interval ceo niz, a kroz rekurzvine pozive se polovi i svaki cvor pokriva manji deo pocetnog niza
void build_segment_tree (std::vector<int> &array, std::vector<int> &segment_tree, int k, int x, int y)
{
  // Ukoliko smo dosli do jednoclanog niza (list u stablu), nema vise polovnjenja intervala pa zavrsavamo rekurziju
  if (x == y) {
    // k je bas pozicija gde treba da smestimo nas element u stablu, a kako su x i y ista pozicija u nizu array mogli smo da uzmemo i array[y], nema razlike
    segment_tree[k] = array[x];
    return;
  }

  // Trazimo sredinu intervala koji smo pokrili na visem nivou, sada imamo njegov levi i desni deo koje pokrivaju potomci cvora sa viseg nivoa. Dakle ako imamo deo niza 1, 2, 3, 4
  // i roditeljski cvor C1 pokriva 1, 2, 3, 4 onda njegovi potomci C2 i C3 pokrivaju redom 1, 2 i 3, 4.
  int middle = (x + y) / 2;

  // Pozivamo rekurziju za kreiranje potomaka cvora (levo i desno)
  // Levi potomak je na poziciji 2 * k u odnosu na poziciju oca
  build_segment_tree(array, segment_tree, 2 * k, x, middle);
  // Desni potomak je na poziciji 2 * k + 1 u odnosu na poziciju oca
  build_segment_tree(array, segment_tree, 2 * k + 1, middle + 1, y);

  // Nakon sto smo kreirali oba sina, potrebno je odrediti koji broj se cuva u ocu, to je zbir brojeva potomaka (podsetiti se strukture stabla)
  // Ukoliko ne bismo imali ovu liniju, samo bi listovi imali neke vrednosti, unutrasnji cvorovi ne bi bili popunjeni na odgovarajuci nacin
  segment_tree[k] = segment_tree[2 * k] + segment_tree[2 * k + 1];
}

int sum_of_segment(std::vector<int> &segment_tree, std::vector<int> &lazy_values, int x, int y, int a, int b, int k)
{
  // Ako postoji neka vrednost u vektoru lazy_values, to znaci da je bilo potrebno vrsit azuriranje segmenta koji pokriva
  // trenutni cvor
  if (lazy_values[k] != 0) {
    // Prvo azuriramo vrednost cvora na odgovarajuci nacin
    // Posto su sve vrednosti iz intervala koji pokriva trenutni
    // cvor uvecane za lazy_values[k], sama vrednost sume se uvecava
    // za lazy_values[k] * broj_elemenata.
    segment_tree[k] += (y - x + 1) * lazy_values[k];

    // Ako se ne nalazimo u listovima, propagiramo azuriranu vrednost
    // na potomke
    if (x != y) {
      lazy_values[2 * k] += lazy_values[k];
      lazy_values[2 * k + 1] += lazy_values[k];
    }

    // Kada smo azurirali vrednost odgovarajuceg cvora, resetujemo vrednost
    // u lenjom stablu
    lazy_values[k] = 0;
  }

  // Ukoliko nema preklapanja segmenata vracamo 0
  if (x > b || y < a)
    return 0;

  // Ukoliko se ceo [x, y] nalazi unutar [a, b] vracamo sumu koju cuva cvor koji pokriva segment [x, y], samim tim i sumu segmenta [x, y]
  if (x >= a && y <= b)
    return segment_tree[k];

  // Delimo segment [x, y] na 2 podsegmenta i problem resavamo rekurzivno za njih
  int middle = (x + y) / 2;

  return sum_of_segment(segment_tree, lazy_values, x, middle, a, b, 2 * k) +
                            sum_of_segment(segment_tree, lazy_values, middle + 1, y, a, b, 2 * k + 1);
}

void update_tree(std::vector<int> &segment_tree, std::vector<int> &lazy_values, int k, int x, int y, int a, int b, int diff)
{
  // Ako postoji neka vrednost u vektoru lazy_values, to znaci da je bilo potrebno vrsit azuriranje segmenta koji pokriva
  // trenutni cvor
  if (lazy_values[k] != 0) {
    // Prvo azuriramo vrednost cvora na odgovarajuci nacin
    // Posto su sve vrednosti iz intervala koji pokriva trenutni
    // cvor uvecane za lazy_values[k], sama vrednost sume se uvecava
    // za lazy_values[k] * broj_elemenata.
    segment_tree[k] += (y - x + 1) * lazy_values[k];

    // Ako se ne nalazimo u listovima, propagiramo azuriranu vrednost
    // na potomke
    if (x != y) {
      lazy_values[2 * k] += lazy_values[k];
      lazy_values[2 * k + 1] += lazy_values[k];
    }

    // Kada smo azurirali vrednost odgovarajuceg cvora, resetujemo vrednost
    // u lenjom stablu
    lazy_values[k] = 0;
  }

  if (x > b || y < a)
    return ;

  // Ukoliko se ceo [x, y] nalazi unutar [a, b] vracamo sumu koju cuva cvor koji pokriva segment [x, y], samim tim i sumu segmenta [x, y]
  if (x >= a && y <= b) {
    segment_tree[k] += (y - x + 1) * diff;

    if (x != y) {
      lazy_values[2 * k] += diff;
      lazy_values[2 * k + 1] += diff;
    }

    return ;
  }

  // Delimo segment na 2 dela
  int middle = (x + y) / 2;

  update_tree(segment_tree, lazy_values, 2 * k, x, middle, a, b, diff);
  update_tree(segment_tree, lazy_values, 2 * k + 1, middle + 1, y, a, b, diff);

  // Nakon azuriranja lista trebalo bi azurirati i sve cvorove koji su na putu od korena do njega jer su se automatski i njihove vrednosti izmenile
  segment_tree[k] = segment_tree[2 * k] + segment_tree[2 * k + 1];
}

int main ()
{
  std::vector<int> array = { 1, 2, 3, 4, 5, 6, 7, 8 };

  int n = array.size();

  int height = ceil(log2(n));

  int size = 2 * pow(2, height);

  array.resize(pow(2, height));
  std::vector<int> segment_tree(size);

  std::vector<int> lazy_values(size, 0);

  build_segment_tree(array, segment_tree, 1, 0, pow(2, height) - 1);

  update_tree(segment_tree, lazy_values, 1, 0, n - 1, 2, 5, 4);

  update_tree(segment_tree, lazy_values, 1, 0, n - 1, 1, 3, 7 );

  std::cout << sum_of_segment(segment_tree, lazy_values, 0, n - 1, 2, 4, 1) << "\n";

  return 0;
}