Алгоритмы и структуры данных на C++: деревья отрезков

Интересуетесь "плюсами"? В статье рассмотрим фундаментальные вещи, такие как алгоритмы и структуры данных в C++. Говорим о деревьях отрезков.

Вкратце о том, что из себя представляет задача на запрос по диапазону: дан массив a из n элементов, дано кол-во запросов q. Для каждого запроса даны границы l, r такие, что 1 ≤ l ≤ r ≤ n.

Требуется для каждого запроса найти максимум/минимум/сумму/НОД и проч. среди элементов a(l), a(l + 1) … a(r - 1), a(r).

Предположим, мы имеем массив длины n = 105 и кол-во запросов q = 105. Если искать максимум/сумму и проч. в массиве «наивным» методом, асимптотическая сложность которого O(n), в задаче на обработку запросов мы будем иметь сложность O(nq), что в данном случае – 1010. При том, что современные ЭВМ обрабатывают 108 о/с, данное решение не является оптимальным.

Давайте построим такую структуру данных, которая сможет обрабатывать запросы на ассоциативные операции за время лучшее, чем линейное, например, логарифмическое. Стоит отметить, что дерево отрезков можно применять тогда и только тогда, когда помимо ассоциативных операций присутствует изменение элементов. Структуры данных, способные обрабатывать запросы без изменения (sparse table, массив префиксных сумм, дерево Фенвика) являются более простыми и будут рассмотрены в следующих статьях.

Итак, что же из себя представляет дерево отрезков? Пусть дан следующий массив: [3, 6, 9, 2, 2, 1]. Корень дерева – сумма всех элементов [0..5] (23). Делим элементы пополам и составляем сумму детей корня – [0..2] (18) и [3..5] (5) соответственно. Далее каждую из этих сумм делим опять пополам – [0..1] (9), [2..2] (9), [3..4] (4), [5..5] (1). Те суммы, у которых левая граница не равна правой, делим ещё раз, получаем соответствующие элементы массива. Вот, как это выглядит:

Иллюстрация дерева отрезков. Жёлтым выделены листья (конечные элементы).

Обратите внимание, что у каждого элемента есть своё отдельное место, и никакая пара элементов не пересекается, поэтому все элементы мы можем записать в массив: рассматриваемая вершина имеет индекс v, а его потомки – индексы 2*v + 1 и 2*v + 2 для левого и правого потомка соответственно.

Расход памяти дерева отрезков будет составлять 4n (доказательство).

Пусть у нас есть массив tree[4*n] и есть вершина tree[v], которая должна хранить сумму диапазона [L..R]. Если L == R, то очевидно, что tree[v] = a[L]. В противном случае мы не можем сразу сказать, чему равна сумма. Для этого нам нужно передать запрос суммы потомкам, а затем записать в tree[v] результат сложения сумм потомков. Пусть M = середина диапазона [L..R], тогда левый потомок t[2*v + 1] соответствует диапазону [L..M], а правый t[2*v + 2][M..R].

Напишем рекурсивную функцию построения дерева отрезков для вычисления суммы:

#include <bits/stdc++.h>

using namespace std;

const int SIZE = (1 << 17);
int a[100100]; // исходный массив
int tree[2 * SIZE]; 
int n; // кол-во элементов в исходном массиве

void build_sum_tree(int v, int L, int R)
{
	if (L == R - 1) // Условие выхода
	{
		if (L < n) // Поскольку мы объявляем большую размерность, необходимо следить за границей
		{
			tree[v] = a[L];
		}
		return; // Присвоили, возвращаемся
	}
	int M = (L + R) / 2; // Выбираем середину отрезка [L..R]
	build_sum_tree(2 * v + 1, L, M); // Запускаем сумму для левого потомка
	build_sum_tree(2 * v + 2, M, R); // И для правого
	tree[v] = tree[2 * v + 1] + tree[2 * v + 2]; // Обновляем текущую вершину
}

// Usage:
// build_sum_tree(0, 0, SIZE); Строим дерево от нулевой вершины с границей от 0 до SIZE (можно n)

Для вычисления max/min/gcd построение аналогично, только tree[v] надо присваивать max(л.п., п.п.)/min(л.п., п.п.)/gcd(л.п., п.п.) соответственно.

Подсчёт суммы на отрезке

Пусть дан массив a, q запросов с данными l, r. Необходимо для каждого запроса найти сумму на отрезке [l..r]. Поскольку в данной реализации мы всё индексируем с нуля, l мы рассматриваем включительно (предварительно уменьшив l на единицу), а r – не включительно.

Утверждается, что в рекурсивной функции подсчёта суммы может быть обработано три ситуации.

  1. [l..r] не пересекается с [L..R] → возвращаем нейтральное значение для суммы;
  2. [L..R] ∈ [l..r] → возвращаем значение текущей вершины;
  3. Третий случай, когда ни один из предыдущих не работает. Нужно получить суммы от потомков текущей вершины. Запрос аналогичен с запросом в функции build_sum_tree. Мы должны получить два слагаемых от двух потомков соответственно и вернуть их сумму. При этом стоит заметить, что рекурсия не будет разрастаться: каждый запрос будет обрабатываться за 2*logN, что асимптотически верно O(logN), поскольку одна из двух вызываемых рекурсий всегда будет возвращать либо 0, либо tree[v]. В этом несложно убедиться, если попробовать пройтись алгоритмом по дереву самостоятельно.

Напишем эту функцию:

int64_t get_summary(int v, int L, int R, int l, int r) // L, R - обрабатываемые функцией границы, l, r - границы запроса
{
    if (r <= L || R <= l) return 0; // Первое условие
    if (l <= L && R <= r) return tree[v]; // Второе условие
    int M = (L + R) / 2;
    int64_t first_child = get_summary(2 * v + 1, L, M, l, r);
    int64_t second_child = get_summary(2 * v + 2, M, R, l, r);
    return first_child + second_child;
}

// Usage:
// build_sum_tree(0, 0, SIZE);
// int l, r; cin >> l >> r; --l;
// int sum = (get_summary(0, 0, SIZE, l, r));

Изменение элемента

Смысл алгоритма донельзя прост: нужно рекурсивно спуститься до нужного нам элемента, затем пересчитать его предков. Выход из рекурсии – если L == R. Тут важно не забыть поменять не только значение элемента в дереве, но ещё и значение в самом массиве. Пусть у нас дан индекс i элемента, который мы хотим заменить, и новое значение x. Если выход из рекурсии не выполняется, то, как всегда, ищем медиану. Если данный индекс меньше медианы, то выполняем запрос изменения элемента к левому потомку, иначе – к правому. Затем пересчитываем tree[v].

void set_in_sum_tree(int v, int L, int R, int i, int x)
{
    if (L == R - 1)
    {
        tree[v] = x;
        a[i] = x;
        return;
    }
    int M = (L + R) / 2;
    if (i < M) set_in_sum_tree(2 * v + 1, L, M, i, x);
    else set_in_sum_tree(2 * v + 2, M, R, i, x);
    tree[v] = tree[2 * v + 1] + tree[2 * v + 2];
}

// Usage:
// build_sum_tree(0, 0, SIZE);
// int i, x; cin >> i >> x; --i;
// set_in_sum_tree(0, 0, SIZE, i, x);

Изменения на отрезке

Имеем всё тот же массив a, q запросов, дерево суммы tree. Необходимо создать изменение элементов (прибавить x) на запрашиваемом отрезке [l..r].

Создадим вспомогательный массив tree_add[2 * SIZE]. Если в tree[v] хранилось значение суммы на отрезке [L..R], то в tree_add[v] будет хранится значение, добавляемое всем элементам на отрезке [L..R]. Тогда нам придётся переписать некоторые методы, поскольку сумма на отрезке будет вычисляться по другой формуле, а именно:

tree[v] + tree_add[v] * (R - L).

Но есть ещё одна проблема: когда запрос на изменение отрезка приходит в определённый отрезок [L..R], он может не дойти ниже (см. иллюстрацию в начале статьи). Представьте, что дан запрос на изменение отрезка [3..5]. Тогда, исходя из нашей логики, отрезок [3..5] может быть "проинформирован" о том, что его сумма поменялась, а отрезки [3..4], [5..5] и прочие – нет.

Эта проблема решается "ленивым проталкиванием изменений" (lazy propogation). Вкратце опишем алгоритм:

  1. Если tree_add[v] != 0, то пересчитываем tree[v];
  2. Передаём tree_add[v] потомкам, если таковые имеются;
  3. Меняем значение tree_add[v] на нейтральное. В случае суммы это ноль.

Напишем функцию push:

void push(int v, int L, int R)
{
    if (L != R - 1) // Условие проталкивания
    {
        tree_add[2 * v + 1] = tree_add[v];
        tree_add[2 * v + 2] = tree_add[v];
    }
    tree[v] += tree_add[v] * (R - L); // Пересчёт
    tree_add[v] = 0;
}

И функцию изменения на отрезке:

void set_range(int v, int L, int R, int l, int r, int x)
{
    push(v, L, R);
    if (r <= L || R <= l) return;
    if (l <= L && R <= r)
    {
        tree_add[v] += x;
        push(v, L, R);
        return;
    }
    int M = (L + R) / 2;
    set_range(2 * v + 1, L, M, l, r, x);
    set_range(2 * v + 2, M, R, l, r, x);
    tree[v] = tree[2 * v + 1] + tree[2 * v + 2];
}

// Usage:
// build_sum_tree(0, 0, SIZE);
// int l, r, x; cin >> l >> r >> x; --l;
// set_range(0, 0, SIZE, l, r, x);

Как и было сказано ранее, необходимо изменить функцию get_summary, добавив push. Ещё один важный момент: используя изменения на диапазоне, мы не поддерживаем изменения самих элементов массива. 

Итак, вот наш код получившегося дерева отрезков:

#include <bits/stdc++.h>

using namespace std;

const int SIZE = (1 << 17);
int a[100100]; // исходный массив
int tree[2 * SIZE]; 
int tree_add[2 * SIZE];
int n; // кол-во элементов в исходном массиве

void push(int v, int L, int R)
{
    if (L != R - 1) // Условие проталкивания
    {
        tree_add[2 * v + 1] = tree_add[v];
        tree_add[2 * v + 2] = tree_add[v];
    }
    tree[v] += tree_add[v] * (R - L); // Пересчёт
    tree_add[v] = 0;
}

void build_sum_tree(int v, int L, int R)
{
	if (L == R - 1) // Условие выхода
	{
		if (L < n) // Поскольку мы объявляем большую размерность, необходимо следить за границей
		{
			tree[v] = a[L];
		}
		return; // Присвоили, возвращаемся
	}
	int M = (L + R) / 2; // Выбираем серидину отрезка [L..R]
	build_sum_tree(2 * v + 1, L, M); // Запускаем сумму для левого потомка
	build_sum_tree(2 * v + 2, M, R); // И для правого
	tree[v] = tree[2 * v + 1] + tree[2 * v + 2]; // Обновляем текущую вершину
}

// Usage:
// build_sum_tree(0, 0, SIZE); Строим дерево от нулевой вершины с границей от 0 до SIZE (можно n)

int64_t get_summary(int v, int L, int R, int l, int r) // L, R - обрабатываемые функцией границы, l, r - границы запроса
{
    push(v, L, R);
    if (r <= L || R <= l) return 0; // Первое условие
    if (l <= L && R <= r) return tree[v]; // Второе условие
    int M = (L + R) / 2;
    int64_t first_child = get_summary(2 * v + 1, L, M, l, r);
    int64_t second_child = get_summary(2 * v + 2, M, R, l, r);
    return first_child + second_child;
}

// Usage:
// build_sum_tree(0, 0, SIZE);
// int l, r; cin >> l >> r; --l;
// int sum = (get_summary(0, 0, SIZE, l, r));

void set_in_sum_tree(int v, int L, int R, int i, int x)
{
    if (L == R - 1)
    {
        tree[v] = x;
        a[i] = x;
        return;
    }
    int M = (L + R) / 2;
    if (i < M) set_in_sum_tree(2 * v + 1, L, M, i, x);
    else set_in_sum_tree(2 * v + 2, M, R, i, x);
    tree[v] = tree[2 * v + 1] + tree[2 * v + 2];
}

// Usage:
// build_sum_tree(0, 0, SIZE);
// int i, x; cin >> i >> x; --i;
// set_in_sum_tree(0, 0, SIZE, i, x);

void set_range(int v, int L, int R, int l, int r, int x)
{
    push(v, L, R);
    if (r <= L || R <= l) return;
    if (l <= L && R <= r)
    {
        tree_add[v] += x;
        push(v, L, R);
        return;
    }
    int M = (L + R) / 2;
    set_range(2 * v + 1, L, M, l, r, x);
    set_range(2 * v + 2, M, R, l, r, x);
    tree[v] = tree[2 * v + 1] + tree[2 * v + 2];
}

// Usage:
// build_sum_tree(0, 0, SIZE);
// int l, r, x; cin >> l >> r >> x; --l;
// set_range(0, 0, SIZE, l, r, x);

int main()
{
    cin >> n;
    for (int i = 0; i < n; ++i) cin >> a[i];

    build_sum_tree(0, 0, SIZE);

    int q;
    cin >> q;

    for (int req = 0; req < q; ++req)
    {
        string type;
        cin >> type;
        if (type == "1")
        {
            int l, r, x;
            cin >> l >> r >> x;
            --l;
            set_range(0, 0, SIZE, l, r, x);
        }
        else if (type == "2")
        {
            int i, x;
            cin >> i >> x;
            --i;
            set_in_sum_tree(0, 0, SIZE, i, x);
        }
        else if (type == "3")
        {
            int l, r;
            cin >> l >> r;
            --l;
            cout << get_summary(0, 0, SIZE, l, r) << endl;
        }
        else if (type == "4")
        {
            for (int i = 0; i < n; ++i) cout << a[i] << " ";
            cout << endl;
        }
    }
    return 0;
}

Вот и все ;)

Сможете решить несколько задач по теме? Вот ресурсы для практики: первый, второй и третий.

А что еще вы хотели бы узнать о программировании на C++?

ЛУЧШИЕ СТАТЬИ ПО ТЕМЕ

matyushkin
29 марта 2020

ТОП-10 книг по C++: от новичка до профессионала

Книги по C++ на русском языке с лучшими оценками. Расставлены в порядке воз...
admin
21 февраля 2017

Какие алгоритмы нужно знать, чтобы стать хорошим программистом?

Данная статья содержит не только самые распространенные алгоритмы и структу...