Дерево отрезков: решение задач на запросы к диапазонам

Дерево отрезков: ваш ключ к быстрым запросам к диапазонам

Когда вы впервые сталкиваетесь с задачей, где нужно много раз обновлять элементы массива и запрашивать сумму подмассива, наивный подход с циклом по каждому запросу даёт сложность O(n) на операцию. Это как вычерпывать тонущий корабль чайной ложкой. Тесты падают по времени, интервьюер смотрит с недоумением. Но есть способ «предобработать» данные так, чтобы и обновления, и запросы выполнялись быстро. Этот способ — дерево отрезков. В этой статье вы научитесь строить дерево отрезков на Python, выполнять запросы суммы и обновления за O(log n) и избегать типичных ошибок.

Что такое дерево отрезков и как оно работает

Дерево отрезков — это структура данных, которая хранит агрегированные значения (например, сумму) для различных отрезков массива. Каждый узел дерева представляет собой отрезок, а его значение — результат операции (сумма, минимум, максимум и т.д.) для этого отрезка. Благодаря этому любой запрос к диапазону можно разбить на O(log n) предвычисленных отрезков, которые полностью покрывают нужный диапазон.

Идея: разбиение диапазона на подотрезки

Представьте массив как батон хлеба. Дерево отрезков разрезает его пополам, затем каждую половину ещё пополам и так далее, помечая каждый кусок его весом. Чтобы узнать вес произвольного среза, вы берёте наименьшее количество уже взвешенных кусков, которые точно его покрывают, — и не нужно взвешивать весь батон каждый раз. Аналогично, при обновлении одного элемента меняются только узлы на пути от листа к корню — их тоже O(log n).

Почему это работает быстро

Любой диапазон можно покрыть не более чем 2·log₂(n) узлами дерева. Это как навигация по метро: вместо того чтобы идти пешком через каждый квартал, вы используете несколько линий, которые довозят вас почти до места. Дерево даёт вам эти «линии» — узлы, которые вместе точно покрывают запрос, и вы просто комбинируете их значения.

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

Начнём с построения дерева. Мы будем использовать массив размером 4n (или 2·2^ceil(log₂ n)), чтобы гарантировать достаточное место для всех узлов.

def build(arr):
    n = len(arr)
    size = 1
    while size < n:
        size <<= 1  # следующий степень двойки
    tree = [0] * (2 * size)
    # Заполняем листья (индексы size ... size+n-1)
    for i in range(n):
        tree[size + i] = arr[i]
    # Строим внутренние узлы снизу вверх
    for i in range(size - 1, 0, -1):
        tree[i] = tree[2*i] + tree[2*i + 1]
    return tree, size

Здесь листья хранят исходные элементы, а каждый родитель — сумму двух детей. Цикл идёт от последнего внутреннего узла к корню, поэтому при вычислении узла его дети уже готовы — это классическое динамическое программирование на дереве.

Запрос суммы на диапазоне

def query(tree, size, l, r):
    # индексы l и r включительно
    l += size
    r += size
    res = 0
    while l <= r:
        if l % 2 == 1:  # l — правый ребёнок: берём его и сдвигаемся
            res += tree[l]
            l += 1
        if r % 2 == 0:  # r — левый ребёнок: берём его и сдвигаемся
            res += tree[r]
            r -= 1
        l //= 2
        r //= 2
    return res

Мы поднимаемся от листьев к корню. Когда левый указатель является правым ребёнком, этот узел целиком лежит в диапазоне, поэтому мы добавляем его значение и пропускаем его. Аналогично для правого указателя, когда он является левым ребёнком. На каждом уровне мы посещаем не более двух узлов, итого O(log n).

Обновление элемента

def update(tree, size, idx, value):
    pos = size + idx
    tree[pos] = value  # меняем лист
    pos //= 2  # переходим к родителю
    while pos:
        tree[pos] = tree[2*pos] + tree[2*pos + 1]
        pos //= 2

Здесь мы обновляем лист и затем пересчитываем всех предков до корня. Сложность — O(log n).

Типичные ошибки и как их избежать

• Ошибка на единицу в цикле запроса: вы либо пропускаете элемент, либо включаете лишний. Проверяйте условие цикла l <= r и не забывайте про сдвиг индексов на size.
• Забыли пересчитать родителей после обновления: устаревшие суммы искажают будущие запросы. Всегда поднимайтесь до корня и обновляйте каждого родителя.
• Размер, не являющийся степенью двойки: формулы для детей (2*i, 2*i+1) ломаются. Дополните массив до следующей степени двойки (или используйте более общую реализацию).

Практическое применение и задачи

С деревом отрезков вы можете решать популярные задачи на собеседованиях:

• Range Sum Query Mutable (LeetCode 307) — поддержка обновлений и запросов суммы за O(log n).
• Range Minimum Query — замените + на min в построении и запросе, и получите O(log n) минимума на отрезке.

Красота в том, что этот же каркас работает для любой ассоциативной операции: сумма, произведение, НОД, побитовое ИЛИ, даже для пользовательских структур (например, хранить сумму и максимум для ответа на «максимальную сумму подмассива»). Как только вы поймёте суть — разбиение диапазона на O(log n) предвычисленных кусков, — вы сможете адаптировать дерево под любую задачу.

Ваш следующий шаг

Возьмите простой массив, постройте дерево отрезков для суммы, выполните несколько смешанных запросов обновления и суммы. Затем замените операцию на произведение или максимум и посмотрите, как тот же код подстраивается. Если застрянете — проследите путь от листа к корню: скорее всего, вы пропустили обновление родителя или неправильно вычислили индекс.

Какую следующую задачу на диапазоны вы покорите с помощью этого инструмента? Поделитесь в комментариях своим кодом или мыслями — интересно увидеть, что вы создадите!