Префиксные суммы: сумма на отрезке за O(1) и подсчёт подмассивов

Короткий ответ. Префиксные суммы это массив, где pref[i] хранит сумму первых i элементов. После одного прохода сумма любого отрезка считается за O(1): pref[r + 1] - pref[l]. Приём нужен, когда запросов на сумму много, а массив не меняется. С хеш-таблицей он ещё и считает подмассивы с заданной суммой. Ниже устройство, код на Python и разбор задачи Subarray Sum Equals K.

Что такое префиксные суммы?

Берём массив a. Строим массив pref длины n + 1. pref[0] равен 0, pref[i] равен сумме a[0] .. a[i - 1].

Сумма отрезка от l до r включительно это pref[r + 1] - pref[l]. Всё, что до l, вычитается. Остаётся ровно отрезок.

Построение стоит O(n). Каждый запрос после этого стоит O(1). Без префиксов каждый запрос стоил бы O(n). При тысяче запросов разница уже в тысячу раз.

По каким словам в задаче узнать префиксные суммы?

• «Много запросов вида сумма на отрезке».
• «Сколько подмассивов с суммой k».
• «Индекс, где сумма слева равна сумме справа».
• «Самый длинный подмассив с суммой 0» или «с равным числом нулей и единиц».

Общий признак: сумма отрезка нужна много раз, а массив не меняется. Если массив меняется между запросами, нужно дерево Фенвика.

Как это выглядит на примере?

Массив [2, 4, 1, 3]. pref = [0, 2, 6, 7, 10]. Сумма от индекса 1 до 2, то есть 4 + 1: pref[3] - pref[1] = 7 - 2 = 5.

Сумма всего массива: pref[4] - pref[0] = 10. Сумма пустого отрезка от 2 до 1: pref[2] - pref[2] = 0. Лишняя нулевая ячейка в начале убирает особый случай для l = 0.

Как написать префиксные суммы на Python?

from itertools import accumulate

def build(a):
    pref = [0]
    for x in a:
        pref.append(pref[-1] + x)
    return pref                      # или list(accumulate(a, initial=0))

def range_sum(pref, l, r):
    return pref[r + 1] - pref[l]     # сумма a[l..r]

def subarray_sum_equals_k(nums, k):
    seen = {0: 1}                    # префикс 0 встречен один раз
    total = count = 0
    for x in nums:
        total += x
        count += seen.get(total - k, 0)
        seen[total] = seen.get(total, 0) + 1
    return count

Третья функция не хранит массив pref. Ей хватает текущей суммы и словаря, сколько раз каждая сумма встречалась.

Как префиксные суммы решают Subarray Sum Equals K?

Нужно число подмассивов с суммой k. Подмассив от l до r имеет сумму pref[r + 1] - pref[l]. Она равна k, когда pref[l] равен pref[r + 1] - k.

Идём слева направо. На каждом шаге смотрим, сколько раз раньше встречался префикс total - k. Столько подмассивов заканчивается здесь. Потом записываем текущий префикс в словарь.

Запись {0: 1} в начале учитывает подмассивы с самого начала. Без неё ответ будет меньше на число таких подмассивов.

Чем префиксные суммы отличаются от соседних приёмов?

Приём	Массив меняется	Запрос	Отрицательные числа
Префиксные суммы	Нет	O(1)	Да
Скользящее окно	Нет	O(n) на весь проход	Нет, условие не монотонно
Дерево Фенвика	Да	O(log n)	Да
Массив разностей	Много прибавлений на отрезках	O(1) на обновление	Да

Как перенести приём на матрицу?

Двумерный префикс pref[i][j] хранит сумму прямоугольника от угла (0, 0) до (i - 1, j - 1). Строится за O(n · m).

Сумма прямоугольника считается по формуле включений и исключений. Вычитаем два прямоугольника выше и левее, потом возвращаем пересечение, которое вычли дважды. Четыре чтения, O(1).

Как решить задачу про подмассив с равным числом нулей и единиц?

Замените каждый ноль на минус один. Теперь нужен подмассив с суммой ноль. Идём с текущей суммой и словарём, где сумма встретилась впервые. Повторилась сумма, значит между двумя встречами лежит подмассив с нулевой суммой. Длина это разница индексов.

На каких ошибках ломается решение?

• Массив pref длины n вместо n + 1. Сумма с нулевого индекса требует особого случая.
• Записывают текущий префикс в словарь до проверки. Подмассив нулевой длины начинает считаться.
• Забывают {0: 1} и теряют подмассивы с начала.
• Берут префиксы там, где массив меняется. Каждое изменение стоит O(n) на пересчёт.

Как объяснить приём на собеседовании?

Скажите: много запросов суммы на отрезке, массив не меняется. Назовите pref и формулу разности. Назовите O(n) на построение и O(1) на запрос.

Если задача про число подмассивов, добавьте словарь. Скажите: «ищу префикс total - k среди прошлых». Это ключевая фраза.

Отработать узнавание приёма можно на странице префиксные суммы. Соседний приём разобран на странице массив разностей.

Вопросы и ответы

Можно ли строить префиксы для произведения или XOR?

Да, для любой операции с обратной. XOR обратен сам себе. Произведение ломается на нулях, там нужен особый учёт.

Как найти индекс, где сумма слева равна сумме справа?

Посчитайте общую сумму total. Идите по массиву, держите сумму слева. Сумма справа это total минус левая минус текущий элемент. Первое совпадение и есть ответ. Один проход, O(1) памяти.

Что делать, если массив меняется?

Пересчитывать префиксы после каждого изменения слишком дорого. Берите дерево Фенвика: O(log n) на изменение и на запрос.

Где приём встречается в реальных системах?

В аналитике: сумма продаж за любой период из накопленных итогов. В графике: интегральные изображения для быстрого размытия. В базах данных: агрегаты по диапазонам.