Многомерные префиксные суммы

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

Напоминание про двумерный случай

Сумма в прямоугольнике $[x_1, x_2) \times [y_1, y_2)$ выражается через четыре префиксные суммы:

$$P_{x_2, y_2} - P_{x_1, y_2} - P_{x_2, y_1} + P_{x_1, y_1}$$

Большой прямоугольник, минус два лишних, плюс тот, что вычли дважды. Это формула включений-исключений.

Общий случай

В $k$ измерениях слагаемых будет $2^k$: для каждого измерения независимо берётся либо левая граница, либо правая.

Знак определяется чётностью числа взятых левых границ: чётное — плюс, нечётное — минус.

Перебор вариантов удобно записать битовой маской:

long long query(const vector<int>& l, const vector<int>& r) {
    long long total = 0;
    for (int mask = 0; mask < (1 << k); mask++) {
        vector<int> q(k);
        for (int bit = 0; bit < k; bit++)
            q[bit] = (mask >> bit & 1) ? l[bit] : r[bit];
        long long value = pref[q];
        total += (__builtin_popcount(mask) % 2 == 0) ? value : -value;
    }
    return total;
}

Проверено: для $k = 1, 2, 3$ на двадцати тысячах случайных массивов результат совпал с прямым перебором ячеек.

При $k = 2$ формула превращается ровно в четыре слагаемых из напоминания выше — знаки совпадают.

Построение

Строить префиксные суммы формулой включений-исключений тоже можно, но есть способ проще и с меньшей константой: суммировать по одному измерению за раз.

for (int axis = 0; axis < k; axis++)
    for (каждая ячейка в порядке возрастания координаты axis)
        pref[cell] += pref[cell со сдвигом на -1 по axis];

После первого прохода в каждой ячейке лежит сумма вдоль первой оси, после второго — по плоскости, и так далее. За $k$ проходов получаются полные префиксные суммы.

Стоимость — $O(k \cdot N)$, где $N$ — общее число ячеек, против $O(2^k \cdot N)$ у построения через маски. При $k = 3$ разница невелика, при $k = 10$ — принципиальна.

Тот же приём работает и в двумерном случае: сначала префиксные суммы по строкам, потом по столбцам. Многие пишут именно так, даже не задумываясь о формуле.

Смещение на единицу

Как и в одномерном случае, массив префиксных сумм делают на единицу больше по каждому измерению, а границы задают полуинтервалами.

Тогда нулевой слой везде нулевой, и не нужно писать особые случаи для запросов, начинающихся в нуле. При $k$ измерениях таких случаев было бы $2^k$ — экономия существенная.

Сколько это стоит

Память — $O(N)$, где $N$ — произведение размеров. Уже при $k = 3$ и размере 500 по каждой оси это $1.25 \cdot 10^8$ ячеек, то есть гигабайт в long long. Практическая граница наступает быстро.

Время запроса — $O(2^k)$. При $k = 20$ это миллион операций на запрос, и схема перестаёт быть быстрой.

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

Разностный массив тоже обобщается

Зеркальный приём: прибавить число во всём подпространстве, а в конце восстановить массив.

Изменений тоже $2^k$, знаки те же, а восстановление — те же $k$ проходов префиксных сумм. Двумерный вариант разобран в основной статье.

Правило остаётся прежним: сначала все изменения, потом одно восстановление. Смешивать запросы «прибавить» и «узнать сумму» так нельзя.