Prefix Sums

Сумма на отрезке считается проходом по нему — это нормально ровно один раз. Когда запросов сто тысяч, а массив на миллион элементов, повторное сложение одних и тех же чисел становится главным расходом: $10^{11}$ операций там, где хватило бы миллиона.

Вторая форма той же боли — «сколько подотрезков имеют сумму k». Перебор пар границ даёт $O(n^2)$ отрезков, и внутри каждого снова складываются те же элементы.

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

индекс:    0   1   2   3   4
a:         3   5   2   7   4

P:     0   3   8  10  17  21      P[i] = a[0] + ... + a[i-1]
       0   1   2   3   4   5      индекс в P

сумма a[2..4] = P[5] - P[2] = 21 - 8 = 13
                              (2 + 7 + 4 = 13)

Тот же массив префиксов отвечает и на вопрос о числе подотрезков с суммой $k$. Условие $P[r]-P[l]=k$ переписывается как $P[l]=P[r]-k$: правая граница у нас в руках, левую достаточно узнать среди уже встреченных.

P = [0] * (len(a) + 1)
for i, x in enumerate(a):
    P[i + 1] = P[i] + x

sum_lr = P[r + 1] - P[l]        # сумма a[l..r] включительно

# сколько подотрезков имеют сумму k
seen = {0: 1}                   # пустой префикс
total = answer = 0
for x in a:
    total += x
    answer += seen.get(total - k, 0)
    seen[total] = seen.get(total, 0) + 1

• Время — $O(n)$ на построение и $O(1)$ на запрос; вариант с хеш-таблицей — $O(n)$ на всё.
• Память — $O(n)$ на массив префиксов, $O(n)$ на таблицу встреченных сумм.
• Предусловие — массив не меняется между запросами.

Сумма телескопируется: $P[r]-P[l]=(a_0+\cdots +a_{r-1})-(a_0+\cdots +a_{l-1})$, и всё, что стоит левее $l$, сокращается. Никаких предположений о знаках элементов при этом не делается — именно поэтому приём работает там, где скользящее окно неверно.

Нулевой элемент $P[0]=0$ — не украшение, а способ убрать частный случай: отрезок, начинающийся с нулевой позиции, считается той же формулой, что и все остальные. В варианте с хеш-таблицей ему соответствует запись seen[0] = 1, и её пропуск теряет ровно те отрезки, что начинаются с начала массива.

• Массив изменился. Одна правка обесценивает весь хвост префиксов; для изменяемых данных нужно дерево Фенвика.
• Забытая запись `seen[0] = 1`. Теряются отрезки, начинающиеся с начала массива, — на большинстве тестов этого не видно.
• Границы. Решите один раз, что означает P[i] — сумму до позиции или включая её, — и держитесь этого во всём коде.
• Переполнение. Сумма миллиона больших чисел выходит за 32-битный тип задолго до конца массива.
• Вещественные числа. Разность двух больших сумм теряет точность; для float лучше суммировать отрезок честно или применять компенсированное суммирование.

1. Почему префиксные суммы работают с отрицательными числами, а скользящее окно — нет?
2. Чему равно P[0] и что сломается, если его не заводить?
3. Сумма делится на k: что класть в хеш-таблицу и почему остаток, а не сама сумма?
4. Массив меняется после каждого десятого запроса — что выгоднее: перестраивать префиксы или взять дерево Фенвика?
5. Как посчитать дисперсию на произвольном отрезке за константу?