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

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

for _ in range(q):
    l, r = map(int, input().split())
    print(sum(a[l - 1:r]))        # проход по отрезку

Оно верное. И оно не пройдёт: каждый запрос стоит до $n$ действий, всего получается $n \cdot q$ — двадцать миллиардов.

Заметьте, что данные при этом не меняются. Все сто тысяч запросов спрашивают об одном и том же списке. Значит можно один раз подготовиться — и отвечать мгновенно.

Это и есть предподсчёт: потратить $O(n)$ один раз, чтобы каждый из $q$ ответов стоил $O(1)$. Вместо $n \cdot q$ получается $n + q$.

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

Построим список, в котором на месте $i$ стоит сумма первых $i$ элементов:

prefix = [0] * (n + 1)
for i in range(n):
    prefix[i + 1] = prefix[i] + a[i]

Обратите внимание на две вещи. Длина массива — $n + 1$, а не $n$. И начинается он с нуля: prefix[0] = 0 означает «сумма пустого начала».

Теперь сумма любого отрезка — это разность двух чисел:

сумма с l-го по r-й = prefix[r] - prefix[l - 1]

Проверим на примере: сумма со второго по четвёртый — это prefix[4] - prefix[1], то есть $21 - 3 = 18$. И правда, $8 + 1 + 9 = 18$.

Именно ради этой формулы массив делают длиннее и начинают с нуля. Если бы prefix[0] был первым элементом, для отрезка, начинающегося с первого, пришлось бы писать особый случай. А так формула работает всегда, включая l = 1: prefix[0] равен нулю, и вычитать нечего.

Массив prefix построен так, что prefix[i] — сумма первых $i$ элементов, а prefix[0] равен нулю.

Чему равна сумма элементов с $l$-го по $r$-й включительно?

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

Нужно узнать, сколько на отрезке положительных чисел? Накапливайте не значения, а единицы и нули:

for i in range(n):
    if a[i] > 0:
        prefix[i + 1] = prefix[i] + 1
    else:
        prefix[i + 1] = prefix[i]

Формула ответа не меняется — разность префиксов теперь даёт количество.

Условие может зависеть и от номера, а не от значения — принцип тот же.

Ограничение одно: признак должен быть известен заранее и не меняться от запроса к запросу. Если бы в каждом запросе спрашивали своё $x$ («сколько больше $x$ на отрезке»), префикс не помог бы — пришлось бы строить его заново на каждый запрос. Такие задачи решаются иначе, и до них мы дойдём.

Пока префиксы отвечали на вопросы «сколько на этом отрезке». Но у формулы есть и обратное прочтение.

Сумма куска с $l$ по $r$ равна $x$ — это то же самое, что

prefix[r] - prefix[l - 1] == x

То есть пара префиксов с разностью $x$. А считать пары с заданной разностью вы умеете с занятия про словари: идти по списку и для каждого нового префикса спрашивать, сколько раз встречалось нужное дополнение.

seen = {0: 1}          # пустой префикс уже встречался один раз
total = 0
count = 0

for value in a:
    total += value
    count += seen.get(total - x, 0)
    seen[total] = seen.get(total, 0) + 1

Строка seen = {0: 1} — не украшение. Без неё потерялись бы куски, начинающиеся с первого элемента: для них нужен «пустой префикс» слева.

Тот же ход решает и задачи про делимость: сумма куска делится на $k$ тогда, когда префиксы на его концах дают одинаковый остаток. Считать надо не сами префиксы, а их остатки.

В списке 10 элементов.

Сколько чисел будет в массиве префиксных сумм, построенном как в занятии? Введите целое число.

Сдвиг на единицу. Самая частая. Признак: ответ верен для всех отрезков, кроме тех, что начинаются с первого элемента, — или наоборот, кроме заканчивающихся последним. Проверяйте на отрезке во весь список и на отрезке из одного элемента.

Забытый нулевой префикс. В задачах про подотрезки без записи seen = {0: 1} теряются все куски, начинающиеся с начала. Признак: ответ стабильно меньше верного.

Медленный ввод. При ста тысячах запросов обычный input() в цикле сам по себе становится заметен. Если решение верное, но не укладывается, читайте всё разом: data = sys.stdin.read().split().

Как проверять себя

• отрезок во весь список — l = 1, r = n;
• отрезок из одного элемента, причём первого и последнего;
• список из одного элемента;
• все нули — в задачах про нулевые суммы это худший случай, и ответ там большой.

Лестница прежняя: пять разминочных, семь основных, три со звёздочкой. Зачёт при десяти решённых.

Разминочные и основные — про запросы к отрезкам: меняется только то, что накапливается. Задачи со звёздочкой — про обратный ход, где префиксы сами становятся данными для словаря.

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

Задача: «дан список из $n$ чисел и $q$ запросов; для каждого выведите сумму элементов с $l$-го по $r$-й, номера с единицы».

Ученик написал:

n, q = map(int, input().split())
a = list(map(int, input().split()))

prefix = [0] * n
prefix[0] = a[0]
for i in range(1, n):
    prefix[i] = prefix[i - 1] + a[i]

for _ in range(q):
    l, r = map(int, input().split())
    print(prefix[r - 1] - prefix[l - 1])

На запросе «со второго по четвёртый» для списка 3 8 1 9 программа выводит 10, хотя верный ответ 18. А на запросе «с первого по четвёртый» выводит 18 вместо 21.

Объясните, что не так, почему ошибка выглядит именно так, и как её исправить.