Дерево отрезков

Обновлено: 2023-12-23
5 мин

Для изучения темы “дерево отрезков” необходимо знать следующие понятия:

  • массивы
  • циклы
  • условные операторы
  • битовые операции

Дерево отрезков (Segment Tree) - это динамическая структура данных, используемая для выполнения операций над интервалами и их обновления. Оно поддерживает две операции: обновление элементов (update) в заданном диапазоне и запрос (query) на сумму элементов в заданном диапазоне.

Выполним следующую задачу: у нас есть массив, и мы хотим находить сумму элементов в определенном диапазоне.

Для этой задачи мы можем использовать дерево отрезков. Оно построено как бинарное дерево, где каждый узел представляет интервал, а значение узла - это сумма элементов в этом интервале.

Основы:

  1. Определение элемента суммы в дереве отрезков:

Элемент суммы в дереве отрезков является суммой всех элементов в диапазоне, который он представляет.

  1. Конструирование дерева отрезков:

Дерево отрезков может быть построено на базе массива чисел. Каждый узел дерева представляет диапазон элементов в массиве и хранит сумму элементов в этом диапазоне.

  1. Реализация операций:

Реализация различных операций в дереве отрезков по сути зависит от его структуры. Однако, существует несколько операций, которые часто используются в различных задачах:

  • Обновление значения в массиве: Эта операция позволяет изменять значение элемента в массиве. Обычно она реализуется с помощью рекурсивного прохода по дереву.

  • Запрос на значение: Эта операция позволяет запрашивать значение элемента в массиве. Обычно она также реализуется с помощью рекурсивного прохода по дереву.

  • Запрос на сумму: Эта операция позволяет запрашивать сумму значений в массиве на заданном интервале. Она обычно реализуется с помощью рекурсивного прохода по дереву и подсчета суммы

Построение дерева отрезков

Так как дерево бинарное, у каждой вершины будет до двух потомков.

Графически это выглядит следующим образом (для массива из 8 элементов):

segment-tree-structure.ru.png

В самой верхней вершине зафиксирован отрезок от начала массива до последнго элемента.

Слева - лева половина от родителя ([0 1 2 3]). Справа - правая половина ()[4 5 6 7]). И так далее до последнего узла с отрезком из одного элемента.

Возьмем массив a = [1, 3, -2, 8, -7]. На его базе построим дерево отрезкови запишем суммы этих отрезков в каждый узел.

Структура такого дерево выглядит следующим образом:

segment-tree-sum.ru.png

💡 Дерево содержит менее 2n вершин. 2*n-1

Число вершин в худшем случае оценивается суммой $n + \frac{n}{2} + \frac{n}{4} + \frac{n}{8} + \ldots + 1 < 2n$

Отобразим такое дерево как массив:

  1. В таком дереве 9 вершин. Массив будет состоять из 9 элементов.
tree[0] = A[0:4]
tree[1] = A[0:2]
tree[2] = A[3:4]
tree[3] = A[0:1]
tree[4] = A[2:2]
tree[5] = A[3:3]
tree[6] = A[4:4]
tree[7] = A[0:0]
tree[8] = A[1:1]

В данном дереве покрыты все вершины.

Имея такую структуру, в значениях вершин можно хранить различные данные, например, сумму отрезка, наименьшее/наибольшее число или другие агрегированные данные на отрезках.

Реализация дерева отрезков на Python

Инициализация дерева

a = [1, 3, -2, 8, -7]

Так как самыми последними вершинами являются отрезки длиной == 1. То процесс создания начинаем с них, постепенно поднимаясь на уровень выше.

💡 Дерево содержит менее 2n вершин. 💡 Нижняя вершина - длина отрезка равна 1.

def build_tree(array):
  n = len(array)
  tree = [0] * 2 * n # Дерево содержит менее **2n** вершин.

  for i in range(n):
    tree[n + i] = a[i] # самые нижние вершины дерева

  # добавляем родителей
  for i in reversed(range(n)):
    tree[i] = tree[2 * i] + tree[2 * i + 1]
    print(i, tree)

>> array = [1, 3, -2, 8, -7]
>> build_tree(array)

  4 [0, 0, 0, 0, 1, 1, 3, -2, 8, -7]
  3 [0, 0, 0, 1, 1, 1, 3, -2, 8, -7]
  2 [0, 0, 2, 1, 1, 1, 3, -2, 8, -7]
  1 [0, 3, 2, 1, 1, 1, 3, -2, 8, -7]
  0 [3, 3, 2, 1, 1, 1, 3, -2, 8, -7]

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

Функция получает индексы исходного массива.

При создании дерева из исходного массива мы помещали каждый отдельный элемент на новый индекс [n + i].

💡 Поэтому когда функция принимает индекс, сначала мы найдет самый нижний элемент в дереве. Он расположен в новом массиве по индексу [длина_исходного_массива + index]

# подсчет суммы на отрезке
def query_tree(l, r):
  global tree, n
  
  sum = 0
  l += n # индекс текущего элемента
  r += n
  while l <= r:
      if l % 2 == 1: # если индекс нечетный
          sum += tree[l]
          l += 1
      if r % 2 == 0:
          sum += tree[r]
          r -= 1
      l //= 2 # floor division. 8 // 3 = 2
      r //= 2
  return sum

>> a = [1, 3, -2, 8, -7]
>> n = len(a)
>> tree = build_tree(a)

>> query_tree(0, 4) # sum([1, 3, -2, 8, -7])
3
>> query_tree(1, 3) # sum([3, -2, 8])
9
>> query_tree(4, 4)
-7

Получаем класс SegmentTree:

Функцию суммирования или любую другую можно включить в момент генерации дерева

class SegmentTree:
    def __init__(self, a):
        self.n = len(a)
        self.tree = [0] * 2 * self.n 
        for i in range(self.n):
            self.tree[self.n + i] = a[i]
        for i in range(self.n - 1, 0, -1):
            self.tree[i] = self.tree[2*i] + self.tree[2*i+1]

    def calculate_sum(self, l, r):
        sum = 0
        l += self.n
        r += self.n
        while l <= r:
            if l % 2 == 1:
                sum += self.tree[l]
                l += 1
            if r % 2 == 0:
                sum += self.tree[r]
                r -= 1
            l //= 2
            r //= 2
        return sum
    
    def find_value(self, l, r):
        l += self.n
        r += self.n
        while l < r:
            if r % 2 == 0:
                r -= 1
            else:
                r -= 1
                l += 1
        return l - self.n

Шаблон класса Segment Tree

class SegmentTree:
    def __init__(self, data, default=0, func=max):
        self._default = default
        self._func = func
        self._len = len(data)
        self._size = _size = 1 << (self._len - 1).bit_length()

        self.data = [default] * (2 * _size)
        self.data[_size:_size + self._len] = data
        for i in reversed(range(_size)):
            self.data[i] = func(self.data[i + i], self.data[i + i + 1])

def __delitem__(self, idx):
    self[idx] = self._default

def __getitem__(self, idx):
    return self.data[idx + self._size]

def __setitem__(self, idx, value):
    idx += self._size
    self.data[idx] = value
    idx >>= 1
    while idx:
        self.data[idx] = self._func(self.data[2 * idx], self.data[2 * idx + 1])
        idx >>= 1

def __len__(self):
    return self._len

def query(self, start, stop):
    """func of data[start, stop)"""
    start += self._size
    stop += self._size
    if start==stop:
        return self._default
    res_left = res_right = self._default
    while start < stop:
        if start & 1:
            res_left = self._func(res_left, self.data[start])
            start += 1
        if stop & 1:
            stop -= 1
            res_right = self._func(self.data[stop], res_right)
        start >>= 1
        stop >>= 1

    return self._func(res_left, res_right)

def __repr__(self):
    return "SegmentTree({0})".format(self.data)

Метод build_tree строит дерево отрезков, а query позволяет выполнять операции запроса.

Ресурсы