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

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

Задача: много сумм на отрезках

Дан список чисел и много вопросов вида «чему равна сумма элементов с l по r?». Можно каждый раз складывать заново:

answer = sum(a[l:r + 1])

Но если в списке 10⁵ чисел и вопросов тоже 10⁵, а отрезки длинные, это около 10¹⁰ сложений — гораздо больше, чем успеет любой компьютер за секунду.

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

Посчитаем заранее массив p, где p[i] — сумма первых i элементов:

a = [5, 2, 7, 1, 3]
p = [0] * (len(a) + 1)
for i in range(len(a)):
    p[i + 1] = p[i] + a[i]
print(p)    # [0, 5, 7, 14, 15, 18]

i	0	1	2	3	4	5
p[i]	0	5	7	14	15	18

p[0] = 0 — сумма нуля элементов. Этот ноль в начале избавляет от особых случаев.

Каждый p[i + 1] получается из предыдущего одним сложением, поэтому весь массив строится за один проход.

Сумма на отрезке за одну операцию

Сумма элементов с индекса l по r включительно (нумерация с нуля):

s = p[r + 1] - p[l]

Почему: p[r + 1] — это сумма элементов a[0]…a[r], а p[l] — сумма a[0]…a[l-1]. Вычитая, убираем лишнее начало.

# сумма a[1] + a[2] + a[3] = 2 + 7 + 1
print(p[4] - p[1])   # 10

Если в задаче элементы нумеруются с единицы, формула выглядит иначе — аккуратно сдвинь индексы. Ошибка на единицу здесь — самая частая.

Быстрый ввод

Когда строк во входе 10⁵ и больше, обычный input() становится заметно медленным. Ускорить можно так:

import sys
input = sys.stdin.readline

После этого input() работает так же, но быстрее. Учти: такой input() оставляет символ перевода строки в конце — для чисел это не важно (int его игнорирует), а строки обрезай через .strip().

Вывод тоже лучше собирать в список и печатать один раз: print("\n".join(map(str, answers))).

Разностный массив

Обратная задача: много раз «прибавить c ко всем элементам с l по r», а в конце узнать все значения. Прибавлять в цикле — снова долго.

Трюк: заводим массив d из нулей и для каждой операции делаем две записи — в начале отрезка и сразу после его конца:

n = 6
d = [0] * (n + 1)

# прибавить 4 к элементам с индекса 1 по 3
d[1] += 4
d[4] -= 4

cur = 0
for i in range(n):
    cur += d[i]
    print(cur, end=" ")   # 0 4 4 4 0 0

Накопленная сумма cur «включает» прибавку в начале отрезка и «выключает» её сразу после конца. Любое число операций обрабатывается за одну запись на операцию плюс один проход в конце.

Префиксные суммы и разностный массив — взаимно обратные: если построить префиксные суммы от разностного массива, получатся сами значения.

Двумерный случай

Префиксные суммы бывают и для таблиц: p[i][j] — сумма прямоугольника от левого верхнего угла до клетки (i - 1, j - 1). Сумму любого прямоугольника тогда можно найти за четыре обращения к p. Это пригодится в более сложных задачах.

Итог

• p[i] — сумма первых i элементов, p[0] = 0, строится за один проход;
• сумма на отрезке [l, r] (с нуля) — p[r + 1] - p[l];
• разностный массив: d[l] += c, d[r + 1] -= c, затем накопить сумму;
• при большом вводе используй sys.stdin.readline.