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


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.

1. Идея

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

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

Поэтому надо один раз сделать предварительные вычисления, а потом отвечать на каждый запрос очень быстро. Для этого используются двумерные префиксные суммы.

Смысл такой:

  • pref[i][j] — сумма всех элементов в прямоугольнике от (1, 1) до (i, j) включительно;
  • тогда сумму любого прямоугольника можно получить по формуле включения-исключения за O(1).

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

Наблюдение 1. Как хранить сумму префикса

Пусть pref[i][j] — сумма элементов в строках от 1 до i и в столбцах от 1 до j.

Тогда:

  • сверху есть сумма pref[i - 1][j],
  • слева есть сумма pref[i][j - 1],
  • но их общая часть учтена дважды, это pref[i - 1][j - 1],
  • и нужно добавить значение текущей клетки.

Получаем формулу:

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

Наблюдение 2. Зачем делать массив на 1 больше

Удобно завести массив pref размера (r + 1) x (c + 1) и заполнить нулём строку 0 и столбец 0.

Тогда:

  • pref[0][j] = 0
  • pref[i][0] = 0

Это избавляет от лишних проверок на границы.

Наблюдение 3. Как найти сумму в произвольном прямоугольнике

Пусть нужно найти сумму в прямоугольнике с углами:

  • верхний левый: (r1, c1)
  • нижний правый: (r2, c2)

Берём сумму большого прямоугольника от (1, 1) до (r2, c2), то есть pref[r2][c2].

Из неё надо убрать:

  • всё, что выше нужного прямоугольника: pref[r1 - 1][c2]
  • всё, что левее нужного прямоугольника: pref[r2][c1 - 1]

Но при этом прямоугольник в верхнем левом углу был вычтен дважды, значит, его надо вернуть:

  • pref[r1 - 1][c1 - 1]

Итоговая формула:

ans = pref[r2][c2] - pref[r1 - 1][c2] - pref[r2][c1 - 1] + pref[r1 - 1][c1 - 1]


3. Алгоритм

  1. Считываем r и c.
  2. Создаём массив pref размера (r + 1) x (c + 1), заполненный нулями.
  3. По мере чтения матрицы сразу строим префиксные суммы:
    • для каждой клетки (i, j) вычисляем pref[i][j] по формуле.
  4. Считываем число запросов q.
  5. Для каждого запроса (r1, c1, r2, c2):
    • вычисляем ответ по формуле включения-исключения;
    • выводим его.

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

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

Построение pref

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

При вычислении:

pref[i - 1][j] содержит сумму всех клеток выше или на строке i - 1, pref[i][j - 1] содержит сумму всех клеток левее или на столбце j - 1.

Если их сложить, область (1..i-1, 1..j-1) попадёт в сумму два раза. Поэтому её надо вычесть один раз, то есть вычесть pref[i - 1][j - 1].

После этого в сумме будут все клетки из прямоугольника (1..i, 1..j), кроме самой клетки (i, j), поэтому добавляем её значение.

Значит, формула действительно правильно строит двумерные префиксные суммы.

Ответ на запрос

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

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

Но прямоугольник (1..r1-1, 1..c1-1) оказался вычтен дважды, поэтому его надо добавить обратно: pref[r1 - 1][c1 - 1].

После этого остаётся ровно сумма нужного прямоугольника.

Следовательно, формула ответа корректна.


5. Сложность

По времени

  • Построение префиксных сумм: O(r * c)
  • Каждый запрос: O(1)
  • Все запросы: O(q)

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

O(r * c + q)

По памяти

Храним массив pref размера (r + 1) x (c + 1).

Итог:

O(r * c)


6. Код на C++17

#include <iostream>
#include <vector>

using namespace std;

int main() {
    int r, c;
    cin >> r >> c;

    vector<vector<long long>> pref(r + 1, vector<long long>(c + 1, 0));

    for (int i = 1; i <= r; i++) {
        for (int j = 1; j <= c; j++) {
            long long x;
            cin >> x;
            pref[i][j] = x + pref[i - 1][j] + pref[i][j - 1] - pref[i - 1][j - 1];
        }
    }

    int q;
    cin >> q;

    for (int k = 0; k < q; k++) {
        int r1, c1, r2, c2;
        cin >> r1 >> c1 >> r2 >> c2;

        long long ans = pref[r2][c2]
                      - pref[r1 - 1][c2]
                      - pref[r2][c1 - 1]
                      + pref[r1 - 1][c1 - 1];

        cout << ans << '\n';
    }

    return 0;
}

7. Код на Python 3

r, c = map(int, input().split())

pref = [[0] * (c + 1) for _ in range(r + 1)]

for i in range(1, r + 1):
    row = list(map(int, input().split()))
    for j in range(1, c + 1):
        x = row[j - 1]
        pref[i][j] = x + pref[i - 1][j] + pref[i][j - 1] - pref[i - 1][j - 1]

q = int(input())

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

Комментарии

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