все Python 3Алгоритмы и структуры данныхпонятие

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

техники

Массив prefix[i] = sum(arr[0..i]). Позволяет за O(1) найти сумму на отрезке [l, r].

Синтаксис

prefix[i] = prefix[i-1] + arr[i]
sum(l, r) = prefix[r] - prefix[l-1]

Главная ловушка — сдвиг индексов: если prefix длиной n+1 и начинается с нуля, то сумма отрезка [l, r] включительно равна prefix[r+1] - prefix[l]; без ведущего нуля случай l == 0 приходится обрабатывать отдельно. Техника окупается только на неизменном массиве: правка одного элемента заставляет пересчитывать весь хвост за O(n), и там уже нужно дерево Фенвика или отрезков. itertools.accumulate строит ту же таблицу лениво, но без initial=0 (появился в 3.8) начинает не с нуля, а с arr[0].

Примеры

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, l, r):
    return prefix[r+1] - prefix[l]
arr = [1,2,3,4,5]
p = build_prefix(arr)
print(range_sum(p, 1, 3))  # → 9 (2+3+4)
arr2 = [3,-1,4,1,5,-9,2,6]
p2 = build_prefix(arr2)
print(range_sum(p2, 0, 7))  # → 11 (вся сумма)
print(range_sum(p2, 2, 4))  # → 10 (4+1+5)
# Количество подмасивов с нулевой суммой
from collections import Counter
def count_zero_sum(arr):
    prefix = 0
    counts = Counter({0: 1})
    result = 0
    for x in arr:
        prefix += x
        result += counts[prefix]
        counts[prefix] += 1
    return result
print(count_zero_sum([1,-1,2,-2,3]))  # → 3 (это [1,-1], [2,-2] и [1,-1,2,-2])
# 2D префиксные суммы
matrix = [[1,2,3],[4,5,6],[7,8,9]]
def build_2d_prefix(mat):
    n, m = len(mat), len(mat[0])
    p = [[0]*(m+1) for _ in range(n+1)]
    for i in range(1,n+1):
        for j in range(1,m+1):
            p[i][j] = mat[i-1][j-1]+p[i-1][j]+p[i][j-1]-p[i-1][j-1]
    return p
p2d = build_2d_prefix(matrix)
print(p2d[-1][-1])  # → 45 (сумма всей матрицы)
# Скользящее окно с prefix
arr3 = [1,2,3,4,5,6,7,8]
p3 = build_prefix(arr3)
k = 3
max_sum = max(range_sum(p3,i,i+k-1) for i in range(len(arr3)-k+1))
print(max_sum)  # → 21 (6+7+8)
# Среднее на подмассиве
print(range_sum(p, 0, 4) / 5)  # → 3.0

См. также

Открыть в полном глоссарииdocs.python.orgСообщить о неточности