Префиксные суммы

Обзор

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

На первый взгляд такие задачи кажутся простыми: достаточно пройтись по массиву и сложить нужные элементы. Но если сумму нужно считать на каждом шаге или для многих диапазонов, наивное решение с вложенными циклами быстро превращается в O(n²). На собеседованиях это частый сигнал, что задачу можно ускорить.

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

Идея паттерна

Идея простая: во время прохода по массиву мы поддерживаем накопленную сумму всех элементов слева от текущей позиции. Благодаря этому нам не нужно каждый раз заново пересчитывать сумму вложенным циклом.

Префиксная сумма для индекса i - это сумма всех элементов от начала массива до этого индекса. Например, для массива:

nums = [1, 7, 3, 6, 5, 6];

префиксные суммы будут такими:

prefix = [1, 8, 11, 17, 22, 28];

Если мы знаем такие накопленные суммы, то можем быстро получить сумму любого подмассива. Сумма от l до r равна:

prefix[r] - prefix[l - 1];

Почему это работает?

prefix[r] хранит сумму всех элементов от 0 до r включительно. В этой сумме есть и нужный нам отрезок l..r, и лишние элементы слева от него - 0..l-1.

prefix[l - 1] - это сумма как раз этой левой части. Если вычесть ее из prefix[r], лишнее уходит, остается только сумма на отрезке от l до r.

Например, для отрезка от 1 до 3 в массиве [1, 7, 3, 6, 5, 6]:

prefix[3] - prefix[0] = 17 - 1 = 16

А 16 - это 7 + 3 + 6, то есть сумма nums[1] + nums[2] + nums[3].

Если l = 0, слева нечего вычитать - берем просто prefix[r].

Когда использовать

Префиксные суммы хорошо подходят, когда в условии задачи есть такие признаки:

  1. Нужно часто считать сумму подмассива.

  2. Нужно сравнивать сумму слева и справа от индекса.

  3. Нужно найти отрезок с заданной суммой.

  4. Есть много запросов вида "какая сумма на диапазоне от l до r".

Наивный подход часто выглядит так: для каждого индекса мы заново считаем сумму слева или справа. Это быстро приводит к сложности O(n²). Префиксные суммы позволяют превратить такие решения в O(n) или даже отвечать на отдельные запросы за O(1) после подготовки.

Два варианта реализации

Есть два распространенных способа использовать этот паттерн.

Первый вариант - создать отдельный массив prefix, где prefix[i] хранит сумму элементов от 0 до i. Это удобно, когда нужно много раз отвечать на запросы по разным диапазонам.

Второй вариант - хранить только одну переменную, например leftSum. Такой подход подходит, когда мы один раз идем по массиву и на каждом шаге сравниваем накопленную сумму с другой величиной. Он экономит память и часто выглядит проще.

Пример

В задаче Индекс равновесия нужно найти индекс, где сумма элементов слева равна сумме элементов справа.

Для массива:

nums = [1, 7, 3, 6, 5, 6];

ответом будет индекс 3, потому что:

1 + 7 + 3 = 5 + 6

То есть слева от числа 6 сумма равна 11, и справа тоже 11.

Обобщения паттерна

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

Суффиксные суммы - зеркальная версия префиксных. Мы идем справа налево и храним сумму всех элементов от текущего индекса до конца массива. Для того же массива суффиксные суммы будут такими:

suffix = [28, 27, 20, 17, 11, 6];

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

Префиксные и суффиксные произведения работают по той же схеме, только вместо сложения накапливаем произведение. Например, для массива [2, 3, 4] префиксные произведения будут [2, 6, 24], а произведение на отрезке от l до r можно получить через деление двух накопленных значений. С произведениями нужно быть внимательнее: ноль в массиве обнуляет целые диапазоны, а при делении важно не делить на ноль.

Префиксные и суффиксные максимумы и минимумы - еще один частый вариант. На каждом шаге запоминаем максимум или минимум всех элементов слева или справа от текущей позиции. Так за O(1) можно узнать, есть ли справа элемент больше текущего, или сравнить текущий элемент с лучшим значением в уже пройденной части массива. Этот прием часто встречается в задачах, где решение для позиции зависит от того, что находится "в будущем" массива.

Итоги

Префиксные суммы - один из базовых паттернов для работы с массивами. Он полезен везде, где повторно считаются суммы на отрезках или нужно понимать баланс между частями массива.

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

Войдите чтобы отмечать прогресс