Много запросов на одни данные

Пусть у нас есть журнал ежедневных изменений баланса: [5, −2, 7, −3, 4]. Пять дней, пять чисел. Положительное число означает приход, отрицательное — расход. Первый день: плюс пять. Второй: минус два. Третий: плюс семь. Четвёртый: минус три. Пятый: плюс четыре. Итого за все пять дней баланс изменился на 11.

Теперь представим, что к этому журналу поступает много вопросов одного типа: «какова сумма изменений за дни с такого-то по такой-то?» Например, со второго по четвёртый включительно: это дни с индексами 1, 2, 3, сумма −2 + 7 + (−3) = 2. Или с третьего по пятый: 7 + (−3) + 4 = 8. Или за все пять дней: 11. Или за пустой промежуток: 0.

Самый прямолинейный способ отвечать на такие вопросы — каждый раз пробегать нужный кусок массива и складывать. Работает? Работает. Но если вопросов много, а куски длинные, мы снова и снова проходим по одним и тем же числам. Каждый запрос стоит O(длины отрезка), и суммарно на q запросов уходит O(q · n). Для больших q и n это заметно. Хочется отвечать на такие вопросы за один шаг, не пересчитывая одно и то же.

Идея префиксных сумм

Идея простая: давайте один раз пройдём по массиву и запишем в отдельную таблицу суммы всех префиксов. То есть для каждой позиции i мы сохраним сумму первых i элементов. Такую таблицу удобно строить слева направо, добавляя по одному элементу.

Заведём массив prefix длины n + 1. Первый элемент положим равным нулю: сумма пустого префикса — это ноль. Это важно: именно из-за него формула запроса получится простой. Дальше для каждого i от 0 до n − 1 делаем prefix[i + 1] = prefix[i] + a[i]. То есть prefix[i + 1] — сумма всех элементов a[0], a[1], …, a[i].

Для нашего массива это выглядит так. prefix[0] = 0. prefix[1] = 0 + 5 = 5. prefix[2] = 5 + (−2) = 3. prefix[3] = 3 + 7 = 10. prefix[4] = 10 + (−3) = 7. prefix[5] = 7 + 4 = 11.

Обратите внимание на смысл каждой ячейки. prefix[k] — это сумма первых k элементов, то есть сумма элементов с индексами от 0 до k − 1. Полный массив — это prefix[5] = 11.

Как из таблицы получить сумму отрезка

Теперь ключевой момент. Пусть нам нужна сумма элементов на позициях с l по r − 1 включительно — то есть полуинтервал [l, r). Сумма первых r элементов равна prefix[r]. Сумма первых l элементов равна prefix[l]. Если из первого вычесть второе, общий префикс «сократится», и останется ровно сумма элементов с l по r − 1. Значит, сумма отрезка [l, r) равна prefix[r] − prefix[l].

Проверим на числах. Отрезок [1, 4): с индексов 1, 2, 3, это −2, 7, −3, сумма 2. По формуле: prefix[4] − prefix[1] = 7 − 5 = 2. Сходится.

Отрезок [2, 5): индексы 2, 3, 4, это 7, −3, 4, сумма 8. По формуле: prefix[5] − prefix[2] = 11 − 3 = 8. Верно.

Полный отрезок [0, 5): prefix[5] − prefix[0] = 11 − 0 = 11. Пустой отрезок, когда l = r, даёт prefix[l] − prefix[l] = 0. Всё согласовано.

Заметьте, насколько это экономно. Один раз прошли по массиву, построили таблицу за O(n). И после этого любой запрос — это одно вычитание, O(1). Больше никаких повторных проходов по кускам массива.

Программа BuildPrefix и RangeSum

Соберём это в программу. Функция BuildPrefix принимает массив и возвращает таблицу префиксных сумм. Функция RangeSum берёт таблицу и границы и возвращает сумму отрезка. Контракт RangeSum: 0 ≤ l ≤ r ≤ n. Если условие нарушено, возвращаем ошибку — это ошибка вызывающей стороны, а не значение. Значение null отклоняется проверкой на входе.

Результат — long, потому что при индексируемой длине массива суммы int-ов могут выходить за пределы int, но укладываться в long. Это допущение: мы не обещаем, что любая сумма любых long-ов поместится в long; речь именно про массивы int с реальными размерами.

using System;

int[] changes = { 5, -2, 7, -3, 4 };
long[] prefix = BuildPrefix(changes);

Console.WriteLine(RangeSum(prefix, 0, 5));
Console.WriteLine(RangeSum(prefix, 1, 4));
Console.WriteLine(RangeSum(prefix, 2, 5));
Console.WriteLine(RangeSum(prefix, 2, 2));

long[] BuildPrefix(int[] a)
{
    if (a == null) throw new ArgumentNullException(nameof(a));
    long[] p = new long[a.Length + 1];
    for (int i = 0; i < a.Length; i++)
    {
        p[i + 1] = p[i] + a[i];
    }
    return p;
}

long RangeSum(long[] prefix, int l, int r)
{
    if (prefix == null) throw new ArgumentNullException(nameof(prefix));
    int n = prefix.Length - 1;
    if (l < 0 || r < l || r > n)
        throw new ArgumentOutOfRangeException();
    return prefix[r] - prefix[l];
}

Вывод:

11
2
8
0

Первая строка — сумма всех дней, 11. Вторая — отрезок [1, 4), это 2. Третья — отрезок [2, 5), это 8. Четвёртая — пустой отрезок [2, 2), это 0.

Цена

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

Но у подхода есть цена, которую важно понимать. Как только массив a меняется, таблица prefix становится неверной: она описывает старое состояние. Если элементы меняются часто, придётся либо каждый раз пересобирать таблицу за O(n), либо искать другой способ. Обновление одного элемента массива не отражается на таблице за O(1): изменение в позиции i затронет все префиксы от i + 1 и дальше. Значит, для частых обновлений префиксные суммы подходят плохо, а для частых запросов на неизменных данных — идеально.

Нулевая сумма отрезка

Отдельно разберём задачу, где префиксные суммы помогают увидеть то, что иначе не бросается в глаза. Спрашивается: есть ли в массиве непустой отрезок, сумма которого равна нулю?

Наивный способ — перебрать все отрезки: для каждой пары границ посчитать сумму. Количество отрезков — порядка n², и перебор становится тяжёлым. Хочется обойтись одним проходом.

Обратимся к префиксным суммам. Сумма отрезка [l, r) равна prefix[r] − prefix[l]. Когда она равна нулю? Ровно тогда, когда prefix[r] = prefix[l]. То есть две позиции в таблице префиксов совпали. При этом l < r, потому что отрезок должен быть непустым. Значит, задача сводится к вопросу: встречается ли какое-то значение префиксной суммы дважды, на разных позициях?

Проверим на примере. Массив [3, −1, −2, 5, −5]. Считаем префиксы: prefix[0] = 0. prefix[1] = 3. prefix[2] = 3 + (−1) = 2. prefix[3] = 2 + (−2) = 0. prefix[4] = 0 + 5 = 5. prefix[5] = 5 + (−5) = 0. Видим: prefix[0] = 0, prefix[3] = 0, prefix[5] = 0. Три позиции с одинаковым значением. Между позициями 0 и 3 лежит отрезок [0, 3): сумма 3 + (−1) + (−2) = 0. Между позициями 3 и 5 лежит отрезок [3, 5): сумма 5 + (−5) = 0. Оба отрезка непусты, оба дают ноль.

Программа HasZeroSum

Для проверки «встречается ли значение дважды» удобно воспользоваться множеством. Заведём HashSet

Стоит уточнить: HashSet мы уже изучали, и как он устроен внутри — не наша забота сейчас. Нам достаточно того, что при подходящих условиях операции Add и Contains выполняются быстро, в среднем за O(1). Именно на этом допущении и держится оценка.

using System;
using System.Collections.Generic;

Console.WriteLine(HasZeroSum(new[] { 3, -1, -2, 5, -5 }));
Console.WriteLine(HasZeroSum(new[] { 1, 2 }));
Console.WriteLine(HasZeroSum(new[] { 0 }));
Console.WriteLine(HasZeroSum(Array.Empty<int>()));

bool HasZeroSum(int[] a)
{
    if (a == null) throw new ArgumentNullException(nameof(a));
    var seen = new HashSet<long> { 0 };
    long running = 0;
    foreach (int value in a)
    {
        running += value;
        if (!seen.Add(running))
        {
            return true;
        }
    }
    return false;
}

Вывод:

True
False
True
False

Разберём каждый случай. На [3, −1, −2, 5, −5] — True: мы нашли два совпадающих префикса, значит, нулевой отрезок существует. На [1, 2] — False: префиксы 0, 1, 3 — все разные, нулевого отрезка нет. На [0] — True: единственный элемент равен нулю, и это непустой отрезок с нулевой суммой. Префиксы: 0 (изначально), потом 0 (после прибавления нуля) — совпадение. На пустом массиве — False: нам нужен непустой отрезок, а в пустом массиве его нет.

Обратите внимание на два места, где легко ошибиться. Первое: мы ищем именно непустой отрезок. Поэтому пустой массив даёт False, хотя формально сумма пустого отрезка равна нулю. Второе: начальное значение 0 в множестве нужно для того, чтобы отрезок, начинающийся с нулевого индекса, тоже был обнаружен. Если текущая сумма в какой-то момент снова стала нулём, значит, от начала массива до этой позиции сумма равна нулю — это непустой отрезок (при условии, что мы уже прошли хотя бы один элемент).

Цена программы HasZeroSum

Идём по массиву один раз, на каждом шаге делаем одно прибавление и одну операцию с множеством. При допущении, что операции множества выполняются в среднем за O(1), общее время — O(n). Дополнительная память — O(n): в худшем случае множество накопит n + 1 различных значений префиксных сумм.

И ещё одна мысль, которую стоит проговорить. Мы не перечисляем отрезки и не храним их. Мы работаем со значениями префиксных сумм, а множество лишь отвечает на вопрос «это значение уже встречалось?» Именно в этом и заключается сила приёма: вместо перебора O(n²) отрезков мы проходим массив один раз и запоминаем всего лишь факт появления каждого значения. Префиксная сумма — это инструмент, который превращает задачу про отрезки в задачу про значения в таблице.