Разбор задачи «Горячие окна K x K в тепловой карте»

1. Идея

Нужно посчитать, сколько квадратов размера K x K в таблице имеют сумму элементов не меньше T.

Если для каждого квадрата отдельно суммировать все K * K клеток, получится слишком медленно: квадратов много, и каждый пересчитывать с нуля нельзя.

Поэтому здесь используется стандартный приём — двумерные префиксные суммы. Они позволяют за O(1) находить сумму любого прямоугольника после предварительной обработки таблицы.

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

Наблюдение 1

Количество возможных квадратов K x K равно:

• по строкам: n - K + 1
• по столбцам: m - K + 1

Итого квадратов может быть до примерно 10^6, если n и m близки к 1000.

Наблюдение 2

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

• O(K^2) на один квадрат,
• всего O((n - K + 1)(m - K + 1)K^2).

При больших размерах это слишком много.

Наблюдение 3

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

Тогда сумму прямоугольника с углами:

• верхний левый (x1, y1)
• нижний правый (x2, y2)

можно найти так:

pref[x2][y2] - pref[x1 - 1][y2] - pref[x2][y1 - 1] + pref[x1 - 1][y1 - 1]

Это и даёт быстрый ответ для каждого квадрата.

3. Алгоритм

1. Считываем n, m, K, T.
2. Строим массив двумерных префиксных сумм pref размера (n + 1) x (m + 1).
   • Нулевая строка и нулевой столбец заполняются нулями.
   • Это удобно, потому что не нужно отдельно обрабатывать выход за границы.

3. Для каждой клетки (i, j) вычисляем:

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

4. Перебираем все возможные верхние левые углы квадратов K x K.

5. Для каждого такого квадрата находим его сумму через префиксные суммы.
6. Если сумма не меньше T, увеличиваем ответ.
7. Выводим ответ.

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

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

Построение pref

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

Формула:

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

работает так:

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

Значит, после заполнения массива pref каждая его ячейка действительно содержит нужную сумму.

Получение суммы квадрата

Рассмотрим квадрат с верхним левым углом (i, j) и нижним правым углом (x2, y2), где:

• x2 = i + K - 1
• y2 = j + K - 1

Сумма всех элементов внутри него равна:

• всей сумме от (1, 1) до (x2, y2),
• минус лишняя верхняя часть,
• минус лишняя левая часть,
• плюс их пересечение, которое вычли дважды.

Именно поэтому используется формула:

pref[x2][y2] - pref[i - 1][y2] - pref[x2][j - 1] + pref[i - 1][j - 1]

Она точно даёт сумму нужного квадрата.

Подсчёт ответа

Мы перебираем все возможные положения квадрата K x K ровно один раз. Для каждого такого квадрата правильно вычисляем сумму и проверяем условие sum >= T.

Следовательно, итоговый счётчик и будет равен количеству горячих квадратов.

5. Сложность

Время

• Построение префиксных сумм: O(n * m)
• Перебор всех квадратов: O((n - K + 1) * (m - K + 1))

Итого: O(n * m)

Память

• Массив префиксных сумм размера (n + 1) x (m + 1)

Итого: O(n * m)

6. Код на C++17

#include <iostream>
#include <vector>

using namespace std;

int main() {
    int n, m, K;
    long long T;
    cin >> n >> m >> K >> T;

    vector<vector<long long>> pref(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;
            pref[i][j] = x + pref[i - 1][j] + pref[i][j - 1] - pref[i - 1][j - 1];
        }
    }

    long long answer = 0;

    for (int i = 1; i + K - 1 <= n; i++) {
        for (int j = 1; j + K - 1 <= m; j++) {
            int x2 = i + K - 1;
            int y2 = j + K - 1;
            long long sum = pref[x2][y2] - pref[i - 1][y2] - pref[x2][j - 1] + pref[i - 1][j - 1];
            if (sum >= T) {
                answer++;
            }
        }
    }

    cout << answer << "\n";
    return 0;
}

7. Код на Python 3

n, m, K, T = map(int, input().split())

pref = [[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]
        pref[i][j] = x + pref[i - 1][j] + pref[i][j - 1] - pref[i - 1][j - 1]

answer = 0

for i in range(1, n - K + 2):
    x2 = i + K - 1
    for j in range(1, m - K + 2):
        y2 = j + K - 1
        s = pref[x2][y2] - pref[i - 1][y2] - pref[x2][j - 1] + pref[i - 1][j - 1]
        if s >= T:
            answer += 1

print(answer)