Разностный массив: прибавление на отрезке за O(1)

Коротко

Разностный массив — зеркальная пара к префиксным суммам. Префиксная сумма превращает массив в «накопленный» и позволяет быстро читать сумму на отрезке. Разностный массив делает обратное: позволяет быстро писать на отрезок, а прочитать результат можно одной префиксной суммой в конце.

Задача, с которой всё начинается

Есть массив из $n$ нулей и $q$ операций вида (l, r, x): прибавить x ко всем элементам с индекса l по r включительно. После всех операций нужно вывести массив.

Наивно каждая операция — цикл по отрезку:

def apply_naive(n: int, updates: list[tuple[int, int, int]]) -> list[int]:
    a = [0] * n
    for l, r, x in updates:
        for i in range(l, r + 1):
            a[i] += x
    return a

Это $O(n \cdot q)$. При $n = q = 10^5$ получаем до $10^{10}$ операций — далеко за пределами любого лимита времени.

Идея: храним разности

Посмотрим не на сами значения, а на то, насколько каждый элемент отличается от предыдущего:

a    = [1, 3, 3, 2, 0]
diff = [1, 2, 0, -1, -2]      # diff[i] = a[i] - a[i-1], diff[0] = a[0]

Обратно массив восстанавливается накопленной суммой: $a[i] = diff[0] + diff[1] + \dots + diff[i]$. То есть a — это префиксная сумма от diff.

Теперь главное. Что происходит с разностями, когда мы прибавляем x ко всем элементам отрезка [l, r]?

• Внутри отрезка соседние элементы выросли одинаково — их разность не изменилась.
• На левой границе a[l] вырос на x, а a[l - 1] нет — diff[l] увеличилась на x.
• Сразу за правой границей a[r + 1] не изменился, а a[r] вырос — diff[r + 1] уменьшилась на x.

Обновление целого отрезка свелось к изменению двух чисел:

Можно думать об этом как о сигнале: в точке l мы «включаем» прибавку x, в точке r + 1 «выключаем». Префиксная сумма несёт включённую прибавку слева направо, пока её не погасят.

Реализация

from itertools import accumulate


def apply_updates(n: int, updates: list[tuple[int, int, int]]) -> list[int]:
    diff = [0] * (n + 1)  # лишняя ячейка для r + 1 == n
    for l, r, x in updates:
        diff[l] += x
        diff[r + 1] -= x
    # префиксная сумма превращает разности обратно в значения
    return list(accumulate(diff[:n]))


apply_updates(5, [(1, 3, 2), (0, 2, 1)])  # [1, 3, 3, 2, 0]

Разберём пример по шагам:

Лишняя ячейка в конце нужна, чтобы обновление с r = n - 1 не выходило за границу. Это тот же приём, что и фиктивный ноль в префиксных суммах: одна дополнительная ячейка убирает ветвление.

Если массив не нулевой

Когда начальный массив уже заполнен, сначала строим из него разности, потом применяем обновления и в конце восстанавливаем:

from itertools import accumulate


class DiffArray:
    def __init__(self, nums: list[int]):
        self.diff = [nums[0]] + [nums[i] - nums[i - 1] for i in range(1, len(nums))]
        self.diff.append(0)

    def add(self, l: int, r: int, x: int) -> None:
        self.diff[l] += x
        self.diff[r + 1] -= x

    def result(self) -> list[int]:
        return list(accumulate(self.diff[:-1]))


d = DiffArray([1, 2, 3, 4, 5])
d.add(1, 3, 10)
d.add(0, 0, -1)
d.result()  # [0, 12, 13, 14, 5]

Можно и проще: применить обновления к нулевому разностному массиву, восстановить прибавки и сложить их с исходным массивом поэлементно. Результат тот же.

Сложность

• Каждое обновление — $O(1)$.
• Восстановление — один проход, $O(n)$.
• Итого $O(n + q)$ по времени и $O(n)$ дополнительной памяти.

Важное ограничение: техника работает в режиме «сначала все записи, потом чтение». Если чтения перемешаны с обновлениями, каждое чтение требует пересчёта префиксной суммы, и выигрыш пропадает. Для такого режима есть дерево Фенвика поверх того же разностного массива (прибавление на отрезке и чтение точки за $O(\log n)$) или дерево отрезков с ленивыми обновлениями.

Задача: Corporate Flight Bookings

LeetCode 1109. Есть n рейсов, пронумерованных с 1. Дан список бронирований [first, last, seats]: на каждый рейс с first по last включительно забронировано seats мест. Вернуть массив, где answer[i] — общее число мест на рейсе i + 1.

Это задача про разностный массив в чистом виде. Единственная ловушка — нумерация с единицы:

from itertools import accumulate


def corp_flight_bookings(bookings: list[list[int]], n: int) -> list[int]:
    diff = [0] * (n + 1)
    for first, last, seats in bookings:
        diff[first - 1] += seats  # переводим в индексацию с нуля
        diff[last] -= seats  # (last - 1) + 1
    return list(accumulate(diff[:n]))


corp_flight_bookings([[1, 2, 10], [2, 3, 20], [2, 5, 25]], 5)
# [10, 55, 45, 25, 25]

Время $O(n + q)$. Именно эту задачу удобно давать на собеседовании как первую: наивное решение очевидно, а оптимальное требует одной идеи.

Задача: Car Pooling

LeetCode 1094. Машина вмещает capacity пассажиров и едет только вперёд. Поездки [numPassengers, from, to]: столько пассажиров садятся в точке from и выходят в точке to. Можно ли развезти всех?

Сформулируем иначе: на каждом участке дороги число пассажиров в машине — это сумма поездок, покрывающих этот участок. Нужно проверить, что максимум не превышает вместимость. Это снова прибавление на отрезке, только отрезок полуоткрытый: [from, to). В точке to пассажиры уже вышли, поэтому гасим прибавку прямо в to, а не в to + 1:

def car_pooling(trips: list[list[int]], capacity: int) -> bool:
    diff = [0] * 1001  # координаты по условию от 0 до 1000
    for passengers, start, end in trips:
        diff[start] += passengers
        diff[end] -= passengers  # на end пассажиры уже вышли
    load = 0
    for delta in diff:
        load += delta
        if load > capacity:
            return False
    return True


car_pooling([[2, 1, 5], [3, 3, 7]], 4)  # False: на участке 3–5 едут 5 человек
car_pooling([[2, 1, 5], [3, 3, 7]], 5)  # True
car_pooling([[2, 1, 5], [3, 5, 7]], 3)  # True: в точке 5 одни вышли, другие сели

Последний пример — проверка на правильные границы. Если бы мы гасили прибавку в end + 1, в точке 5 машина считалась бы загруженной на 5 человек, и ответ был бы неверным.

Если координаты большие

Массив на 1001 ячейку работает, потому что координаты маленькие. Если бы они были до $10^9$, разностный массив не поместился бы в память. Тогда используют ту же идею в виде событий: каждая поездка даёт два события (start, +p) и (end, -p), события сортируются, и мы идём по ним, накапливая сумму:

def car_pooling_sweep(trips: list[list[int]], capacity: int) -> bool:
    events = []
    for passengers, start, end in trips:
        events.append((start, passengers))
        events.append((end, -passengers))
    events.sort()  # в одной точке отрицательные события идут раньше
    load = 0
    for _, delta in events:
        load += delta
        if load > capacity:
            return False
    return True

Это $O(q \log q)$ и не зависит от диапазона координат. По сути это разностный массив, в котором хранятся только ненулевые ячейки. На этой же идее решается задача о минимальном числе переговорок — она разобрана в статье про задачи на интервалы.

Я сверил оба варианта Car Pooling на тысячах случайных входов — ответы совпадают.

Двумерный разностный массив

Та же идея переносится на матрицу: прибавить x ко всем клеткам прямоугольника с углами (r1, c1) и (r2, c2). Теперь нужно отметить четыре угла:

def apply_2d(rows: int, cols: int, updates: list[tuple]) -> list[list[int]]:
    diff = [[0] * (cols + 1) for _ in range(rows + 1)]
    for r1, c1, r2, c2, x in updates:
        diff[r1][c1] += x
        diff[r1][c2 + 1] -= x
        diff[r2 + 1][c1] -= x
        diff[r2 + 1][c2 + 1] += x
    # префиксные суммы по строкам, затем по столбцам
    for r in range(rows):
        for c in range(1, cols):
            diff[r][c] += diff[r][c - 1]
    for r in range(1, rows):
        for c in range(cols):
            diff[r][c] += diff[r - 1][c]
    return [row[:cols] for row in diff[:rows]]


apply_2d(3, 4, [(0, 0, 1, 1, 1), (1, 1, 2, 3, 2)])
# [[1, 1, 0, 0],
#  [1, 3, 2, 2],
#  [0, 2, 2, 2]]

Почему четыре угла: после двумерной префиксной суммы значение в (r1, c1) «растекается» вправо и вниз на весь квадрант. Два минуса отрезают лишнее справа и снизу, а плюс в (r2 + 1, c2 + 1) возвращает угол, который вычли дважды. Это та же формула включений-исключений, что и в двумерных префиксных суммах, только применённая наоборот.

Сложность — $O(1)$ на обновление и $O(rows \cdot cols)$ на восстановление. На LeetCode эта техника нужна, например, в задаче 2536 «Increment Submatrices by One».

Типичные ошибки

• Забыть лишнюю ячейку. Без n + 1 обновление с r = n - 1 падает с IndexError.
• Путать закрытые и полуоткрытые отрезки. Для [l, r] гасим в r + 1, для [l, r) — в r. В Car Pooling это ровно разница между верным и неверным ответом.
• Сбиться с индексацией. В Corporate Flight Bookings рейсы с единицы: first - 1 и last, а не first и last + 1.
• Читать до восстановления. diff[i] — не значение элемента, а разность. Пока не посчитана префиксная сумма, массив читать нельзя.
• Применять, когда чтения идут вперемешку с записями. Тогда пересчёт префиксной суммы на каждое чтение даёт те же $O(n \cdot q)$. Нужна структура с логарифмическими операциями.
• В 2D забыть угол +x. Без четвёртой точки нижний правый квадрант уйдёт в минус.

Как распознать задачу на собеседовании

Признаки:

• много операций вида «прибавь к отрезку», «забронируй диапазон», «включи на интервале»;
• ответ нужен один раз, в конце, или нужен максимум по всем точкам;
• спрашивают «сколько интервалов покрывает каждую точку» или «бывает ли момент, когда одновременно больше k».

Хорошая подача ответа: сначала наивное $O(n \cdot q)$, потом наблюдение, что внутри отрезка разности не меняются, затем две точки на обновление и префиксная сумма в конце. Если интервьюер спрашивает про большие координаты — переход к сортировке событий. Если про чтения посреди обновлений — дерево Фенвика или отрезков.

Что запомнить

• Разностный массив хранит a[i] - a[i-1]; исходный массив — его префиксная сумма.
• Прибавление x на [l, r] — это diff[l] += x и diff[r + 1] -= x.
• Итоговая сложность $O(n + q)$, но только если все чтения после всех записей.
• Для полуоткрытых отрезков гасим прибавку в правой границе, а не после неё.
• При больших координатах тот же приём работает через сортировку событий.
• В 2D отмечаем четыре угла прямоугольника и считаем двумерную префиксную сумму.