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

Один предподсчёт — и сумма любого отрезка считается мгновенно.

Префиксная сумма — массив, где элемент i хранит сумму всех элементов исходного массива от начала до i включительно.

Зачем это нужно

Если нужно много раз спрашивать «какова сумма элементов с индекса L по R», наивно это O(n) на каждый запрос. С префиксными суммами после предподсчёта за O(n) любой такой запрос — O(1): сумма[L..R] = prefix[R+1] - prefix[L].

def build_prefix(arr):
    prefix = [0] * (len(arr) + 1)
    for i, x in enumerate(arr):
        prefix[i + 1] = prefix[i] + x
    return prefix

def range_sum(prefix, left, right):
    return prefix[right + 1] - prefix[left]

arr = [3, 1, 4, 1, 5, 9, 2]
p = build_prefix(arr)
print(p)
print(range_sum(p, 1, 3))  # 1 + 4 + 1
print(range_sum(p, 0, 6))  # вся сумма

Вывод:

[0, 3, 4, 8, 9, 14, 23, 25]
6
25

Классическая задача: подмассивы с суммой k

Сколько непрерывных подмассивов имеют сумму ровно k? Идея: если prefix[j] - prefix[i] == k, то подмассив (i, j) подходит. Перепишем: prefix[i] == prefix[j] - k. Идём слева направо и в хеш-таблице считаем, сколько раз встретилась каждая префиксная сумма.

def subarray_sum_count(arr, k):
    count = 0
    running = 0
    seen = {0: 1}
    for x in arr:
        running += x
        count += seen.get(running - k, 0)
        seen[running] = seen.get(running, 0) + 1
    return count

print(subarray_sum_count([1, 1, 1], 2))
print(subarray_sum_count([1, 2, 3], 3))

Вывод:

2
2

Как работает под капотом

Префиксная сумма — это дискретный аналог интеграла: prefix[i+1] - prefix[i] возвращает исходный элемент (аналог производной). Связка с хеш-таблицей превращает «найти отрезок с нужной суммой» в «найти, встречалось ли нужное значение раньше» — а это O(1) на проверку. Инициализация seen = {0: 1} учитывает подмассивы, начинающиеся с самого начала: пустой префикс имеет сумму 0.

Частые ошибки

Сдвиг индексов: prefix длиннее массива на 1, легко промахнуться на единицу.
Забыть положить {0: 1} — потеряются подмассивы от начала.
Обновлять хеш до прибавления текущего — порядок важен, иначе посчитаете пустой подмассив.

Итог

Префиксные суммы дают сумму отрезка за O(1) после O(n) предподсчёта.
В паре с хеш-таблицей решают «подмассив с суммой k» за один проход.
Следите за сдвигом индексов и инициализацией {0: 1}.