Редакция для Сумма на прямоугольнике в матрице


Remember to use this editorial only when stuck, and not to copy-paste code from it. Please be respectful to the problem author and editorialist.
Submitting an official solution before solving the problem yourself is a bannable offence.

Автор: montes332

1. Идея

Нужно отвечать на много запросов суммы на прямоугольнике матрицы.

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

Стандартный способ ускорить такие запросы — построить двумерные префиксные суммы.

Создадим массив p, где p[i][j] — сумма всех элементов в прямоугольнике от (1, 1) до (i, j) включительно.

Тогда сумму на любом прямоугольнике (r1, c1) ... (r2, c2) можно получать за O(1) по формуле включений-исключений.

2. Наблюдения

Наблюдение 1

Пусть известны суммы для всех верхних прямоугольников. Тогда:

p[i][j] = a[i][j] + p[i-1][j] + p[i][j-1] - p[i-1][j-1]

Почему так:

  • p[i-1][j] — сумма всего сверху,
  • p[i][j-1] — сумма всего слева,
  • их общая часть p[i-1][j-1] посчиталась дважды, поэтому её надо вычесть,
  • и добавить текущую клетку a[i][j].

Наблюдение 2

Сумма на прямоугольнике от (r1, c1) до (r2, c2) выражается через префиксы так:

p[r2][c2] - p[r1-1][c2] - p[r2][c1-1] + p[r1-1][c1-1]

Смысл такой:

  • p[r2][c2] содержит сумму всего прямоугольника от (1,1) до (r2,c2),
  • убираем лишнюю часть сверху: p[r1-1][c2],
  • убираем лишнюю часть слева: p[r2][c1-1],
  • но верхний левый угол убрали дважды, поэтому возвращаем его: p[r1-1][c1-1].

Наблюдение 3

Очень удобно хранить массив p размером (n+1) × (m+1) и считать, что нулевая строка и нулевой столбец состоят из нулей.

Тогда не нужны отдельные проверки на границы: формулы работают одинаково для всех i, j, r1, c1.

Наблюдение 4

Значения матрицы могут быть до 10^9 по модулю, а размер матрицы до 1000 × 1000, поэтому суммы могут быть очень большими. Нужен тип long long в C++.

В Python обычный int подходит автоматически.

3. Алгоритм

  1. Считать n, m, q.
  2. Создать массив p размера (n+1) × (m+1), заполненный нулями.
  3. Для каждой клетки матрицы при чтении сразу вычислять p[i][j] по формуле:
    • p[i][j] = x + p[i-1][j] + p[i][j-1] - p[i-1][j-1]
  4. Для каждого запроса (r1, c1, r2, c2) вычислить ответ:
    • p[r2][c2] - p[r1-1][c2] - p[r2][c1-1] + p[r1-1][c1-1]
  5. Вывести ответ.

4. Почему это работает

Докажем корректность.

Корректность построения p

По определению p[i][j] должна хранить сумму всех элементов из прямоугольника (1,1) ... (i,j).

Формула

p[i][j] = a[i][j] + p[i-1][j] + p[i][j-1] - p[i-1][j-1]

действительно даёт эту сумму:

  • p[i-1][j] покрывает все строки от 1 до i-1,
  • p[i][j-1] покрывает все столбцы от 1 до j-1,
  • область (1,1) ... (i-1,j-1) входит в обе суммы, поэтому её надо вычесть один раз,
  • остаётся добавить клетку (i,j).

Значит, после заполнения массива каждое p[i][j] содержит правильную префиксную сумму.

Корректность ответа на запрос

Рассмотрим прямоугольник (r1, c1) ... (r2, c2).

  • p[r2][c2] содержит сумму всего прямоугольника (1,1) ... (r2,c2).
  • В этой сумме есть лишняя часть выше строки r1, то есть (1,1) ... (r1-1,c2). Её вычитаем: p[r1-1][c2].
  • Также есть лишняя часть левее столбца c1, то есть (1,1) ... (r2,c1-1). Её тоже вычитаем: p[r2][c1-1].
  • Но прямоугольник (1,1) ... (r1-1,c1-1) был вычтен дважды, поэтому его надо прибавить обратно: p[r1-1][c1-1].

Итоговая формула оставляет ровно сумму нужного прямоугольника и ничего лишнего.

Следовательно, алгоритм всегда выдаёт правильный ответ на каждый запрос.

5. Сложность

Построение двумерных префиксных сумм занимает O(n * m).

Каждый запрос обрабатывается за O(1).

Итоговая сложность:

O(n * m + q)

Память:

O(n * m)

6. Код на C++17

#include <iostream>
#include <vector>

using namespace std;

int main() {
    int n, m, q;
    cin >> n >> m >> q;

    vector<vector<long long>> p(n + 1, vector<long long>(m + 1, 0));

    for (int i = 1; i <= n; i++) {
        for (int j = 1; j <= m; j++) {
            long long x;
            cin >> x;
            p[i][j] = x + p[i - 1][j] + p[i][j - 1] - p[i - 1][j - 1];
        }
    }

    for (int k = 0; k < q; k++) {
        int r1, c1, r2, c2;
        cin >> r1 >> c1 >> r2 >> c2;
        long long ans = p[r2][c2] - p[r1 - 1][c2] - p[r2][c1 - 1] + p[r1 - 1][c1 - 1];
        cout << ans << '\n';
    }

    return 0;
}

7. Код на Python 3

n, m, q = map(int, input().split())

p = [[0] * (m + 1) for _ in range(n + 1)]

for i in range(1, n + 1):
    row = list(map(int, input().split()))
    for j in range(1, m + 1):
        x = row[j - 1]
        p[i][j] = x + p[i - 1][j] + p[i][j - 1] - p[i - 1][j - 1]

for _ in range(q):
    r1, c1, r2, c2 = map(int, input().split())
    ans = p[r2][c2] - p[r1 - 1][c2] - p[r2][c1 - 1] + p[r1 - 1][c1 - 1]
    print(ans)

Комментарии

Еще нет ни одного комментария.