Діма і масив#2941

Діма і масив (Розбір)

Задача вимагає обробки операцій двох типів:

присвоєння значення одному елементу масиву;
обчислення суми елементів на відрізку $[f; t]$ .

Нам потрібна структура даних, яка швидко виконує обидві операції. Такою структурою є дерево відрізків для сум. Кожна його вершина зберігає суму елементів свого відрізка, причому значення у внутрішній вершині дорівнює сумі значень її двох дітей.

Дерево відрізків дозволяє як обчислювати суму на відрізку, так і змінювати елемент масиву за $O (lo g n)$ O(\log n).

Приклад

Змоделюємо запити з прикладу. Спочатку масив дорівнює $(1, 2, 3)$ (1, 2, 3).

$? 13$ : сума елементів на відрізку $[1; 3]$ дорівнює $1 + 2 + 3 = 6$ 1 + 2 + 3 = 6;
$= 32$ : елементу $a_{3}$ a_3 присвоюється значення $2$ , і масив стає $(1, 2, 2)$ (1, 2, 2);
$? 13$ : сума елементів на відрізку $[1; 3]$ дорівнює $1 + 2 + 2 = 5$ 1 + 2 + 2 = 5.

Подивимось, як змінюється дерево відрізків. На лівому малюнку зображено дерево для початкового масиву, а на правому — після присвоєння $a_{3} = 2$ a_3 = 2. Для кожної вершини показано її відрізок та суму елементів на ньому; змінені вершини виділено.

Відрізок $[1; 3]$ в обох запитах збігається з відрізком кореня, тому відповідь береться безпосередньо з кореня: $6$ до присвоєння та $5$ після нього. Під час присвоєння $a_{3} = 2$ a_3 = 2 перераховуються лише вершини на шляху від листа $[3]$ до кореня: лист $[3]$ отримує значення $2$ , а корінь отримує значення $3 + 2 = 5$ 3 + 2 = 5.

Реалізація алгоритму

Зберігатимемо дерево відрізків у векторі $S e g T r ee$ SegTree довжиною $4 \cdot M A X$ 4 \cdot MAX, де $M A X = 5 \cdot 1 0^{5}$ MAX = 5 \cdot 10^5 — максимальна кількість елементів масиву. Кожна вершина дерева зберігає суму елементів свого відрізка.

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

Функція BuildTree будує дерево відрізків з масиву $a$ . Вона отримує номер поточної вершини $v$ та межі $l p os$ і $r p os$ відрізка, що відповідає цій вершині.

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

Якщо відрізок вершини $v$ складається з одного елемента ( $l p os = r p os$ lpos = rpos), то $v$ — лист дерева. Записуємо в неї значення $a [l p os]$ .

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

Інакше знаходимо середину відрізка $mi d = (l p os + r p os) /2$ mid = (lpos + rpos) / 2.

    int mid = (lpos + rpos) / 2;

Відрізок $[l p os; r p os]$ [lpos; rpos] розбивається на два підвідрізки $[l p os; mi d]$ [lpos; mid] та $[mi d + 1; r p os]$ [mid + 1; rpos]. Рекурсивно будуємо дерево для лівого нащадка $2 v$ та правого нащадка $2 v + 1$ 2v + 1.

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

Значення у вершині $v$ дорівнює сумі значень її нащадків.

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

Функція Update присвоює значення $v a l$ елементу з індексом $p os$ . Вершина $v$ відповідає відрізку $[l p os; r p os]$ [lpos; rpos].

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

Якщо вершина $v$ — лист ( $l p os = r p os$ lpos = rpos), вона відповідає позиції $p os$ . Записуємо в неї значення $v a l$ .

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

Інакше визначаємо, в якій половині відрізка знаходиться позиція $p os$ , і рекурсивно спускаємось у лівого або правого нащадка.

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

Після повернення з рекурсії перераховуємо суму у вершині $v$ . Таким чином оновлюються всі вершини на шляху від листа до кореня.

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

Функція Summa повертає суму елементів на відрізку $[l e f t; r i g h t]$ [left; right]. Вершина $v$ відповідає відрізку $[l p os; r p os]$ [lpos; rpos]. Якщо запитаний відрізок порожній ( $l e f t > r i g h t$ left > right), повертаємо $0$ — нейтральний елемент для додавання.

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

Якщо запитаний відрізок збігається з відрізком вершини, відповідь зберігається в $S e g T r ee [v]$ SegTree[v].

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

Інакше знаходимо середину відрізка $mi d = (l p os + r p os) /2$ mid = (lpos + rpos) / 2.

  int mid = (lpos + rpos) / 2;

Розбиваємо запитаний відрізок на дві частини: $[l e f t; min (r i g h t, mi d)]$ [left; \min(right, mid)] лежить у відрізку лівого нащадка, а $[max (l e f t, mi d + 1); r i g h t]$ [\max(left, mid + 1); right] лежить у відрізку правого нащадка. Одна з частин може бути порожньою. Повертаємо суму результатів для обох нащадків.

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

Основна частина програми. Зчитуємо розмір масиву $n$ , кількість запитів $q$ та елементи масиву у вектор $v$ , починаючи з індексу $1$ . Виклик scanf("\n") пропускає перехід на новий рядок після масиву, щоб наступним зчитаним символом був знак операції.

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

Будуємо дерево відрізків з елементів $v [1 \dots n]$ v[1 \ldots n].

BuildTree(v,1,1,n);

Обробляємо $q$ запитів по черзі. Спочатку зчитуємо символ операції $c$ ; пробіл у форматі "%c " пропускає пробіл після нього. Якщо $c$ = '=', присвоюємо значення $d$ елементу з індексом $i$ . Інакше обчислюємо та виводимо суму на відрізку $[f; t]$ . Символ \n у форматі "%d %d\n" пропускає перехід на новий рядок, щоб наступне зчитування в $c$ отримало знак наступної операції.

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++