Расскажу о том, как решал одну из наиболее интересных задач в разминке Яндекс Алгоритмы 2023 г. Интересной я называю ее потому, что: 1) решал я кратно дольше, чем предыдущие 6 задач из разминки вместе взятые; 2) именно в этой задаче я проникся мощью префиксных сумм, и применением их для двумерных массивов.
И так задача:
Кролики очень любопытны. Они любят изучать геометрию, бегая по грядкам. Наш кролик как раз такой. Сегодня он решил изучить новую фигуру — квадрат.
Кролик бегает по грядке — клеточному полю N × M клеток. В некоторых из них посеяны морковки, в некоторых нет.
Помогите кролику найти сторону квадрата наибольшей площади, заполненного морковками полностью.

Формат ввода
В первой строке даны два натуральных числа N и M ( 1 ≤ N, M ≤ 1000). Далее в N строках расположено по M чисел, разделенных пробелами (число равно 0, если в клетке нет морковки или 1, если есть).
Формат вывода
Выведите одно число — сторону наибольшего квадрата, заполненного морковками.
Ограничение времени | 4 секунды |
Ограничение памяти | 80Mb |
Пример 1
Ввод
4 5
0 | 0 | 0 | 1 | 0 |
0 | 1 | 1 | 1 | 0 |
0 | 0 | 1 | 1 | 0 |
1 | 0 | 1 | 0 | 0 |
Ответ: 2
Пример 2
Ввод
5 5
1 | 1 | 1 | 1 | 1 |
1 | 1 | 1 | 1 | 1 |
1 | 1 | 1 | 1 | 1 |
1 | 1 | 1 | 1 | 1 |
1 | 1 | 1 | 1 | 1 |
Ответ: 5
Пример 3
Ввод
4 4
0 | 0 | 0 | 0 |
0 | 0 | 0 | 0 |
0 | 0 | 0 | 0 |
0 | 0 | 0 | 0 |
Ответ: 0
И так рассмотрим алгоритм наивного(жадного) кролика. А именно берем самый большой квадрат, который влезает в наш прямоугольник. Назовем его «окном». И начиная с верхнего левого угла начинаем смещаться либо вправо на один столбец (если прямоугольник вытянут по горизонтали), либо вниз на одну строку (если прямоугольник вытянут по вертикали). Конечно на входе может быть и квадрат. Но это не суть важно.
Когда наше «окно» добегает до конца, так что, что его правый нижний угол совпадает с правым нижним углом матрицы – мы должны уменьшить площадь окна на 1 и снова начать смещать с верхнего левого угла.

И так до тех пор, пока не окажется, что наше «окно» наткнется на поле, которое полностью состоит из 1. Рано или поздно это произойдет. Вопрос, разумеется, в том, когда это произойдет и сколько пройдет времени. Фрагмент основного кода этого способа приведен ниже:
while size > 0:
sqr_line = line[ln][x1:x2+1]
if any(val == 0 for val in sqr_line) == True:
x2 += 1 # смещаем вправо правую границу квадрата бегунка
if x2 < wdth: # если НЕ достигли правой границы
x1 += 1 # смещаем вправо левую границу квадрата бегунка
ln = y1
else: # если достигли правой границы
x1 = 0 # и возвращаемся влево
x2 = x1 + size - 1 # и возвращаемся влево
y1 += 1 # смещаемся вниз
y2 += 1 # смещаемся вниз
ln = y1
if y2 >= hght: # если после смещений еще достигли и нижней границы и спускаться некуда
size -= 1
x1 = 0
x2 = x1 + size - 1
y1 = 0
y2 = y1 + size - 1
ln = y1
else:
ln += 1
if ln == y2 + 1:
break
И на смену ему приходит турбо кролик и его префиксные суммы.
Здесь я не буду подробно рассказывать о префиксных суммах, т.к. об этом написано и сказано не один раз в подобных статьях. Вместо этого сфокусируемся на применении префиксных сумм для нашей задачи.
Для начала давайте насчитаем префиксные суммы каждого прямоугольника, входящего в изначальную матрицу максимально быстрым способом, который изображен ниже:

Именно таким способом и происходит быстрое насчитывание префиксных сумм. Если мы хотим получить префиксную сумму зеленого квадрата(или прямоугольника), то нам нужно сложить префиксные суммы верхнего (голубого) прямоугольника и левого(фиолетового) прямоугольника. Далее нужно не забыть вычесть сумму желтого прямоугольника, так как он содержится и в верхнем/голубом и в левом/фиолетовом прямоугольниках. Останется прибавить значение красного элемента и все префиксная сумма посчитана. Гениально и мощно!) Код, который это делает очень простой:
for i in range(hght):#насчитываем префиксные суммы по всему прямоугольнику
for j in range(wdth):
up = (0 if i == 0 else pr_sm[i - 1][j]) # прямоугольник сверху, если идем по верхней границе
lf = (0 if j == 0 else pr_sm[i][j - 1]) # прямоугольник слева, если идем по левой границе
df = (0 if (j == 0 or i == 0) else pr_sm[i - 1][
j - 1]) # общее множество, которое нужно вычесть, если идем по верхней или левой границе
pr_sm[i][j] = up + lf - df + vector2D[i][j]
Решение близко, но это еще не все. Использование префиксных сумм в чистом виде ничего нам не даст. Так как в нашем случае префиксные суммы, лежащие в узлах матрицы, всегда показывают сумму прямоугольника от левого верхнего угла до узла матрицы. Неужели нам никак не помогут префиксные суммы и, отказавшись от них, придется придумывать что-то другое. Но на самом деле ничего супер нового нам в этой задаче выдумывать не придется. Если коротко, то формулу успеха можно описать примерно так:
Решение = принцип расчета префикс. сумм + префиксные суммы;
Теперь перейдем к детальному плану.
Мы будем проверять каждый квадрат на предмет наличия в нем 0. Только делать мы это будем путем сравнения его площади с суммой единиц в этом квадрате. Мы знаем длину стороны квадрата, а значит легко найдем площадь.
Остается решить вопрос: Как быстро получить сумму единиц нужного квадрата? Поможет нам в этом тот же принцип быстрого расчета префиксных сумм. Сам принцип расчета изложен на рисунке ниже:

То есть, чтобы найти сумму зеленого квадрата мы отнимем от префиксной суммы красного квадрата префиксные суммы голубого и фиолетового прямоугольников. А также не забудем прибавить префиксную сумму желтого квадрата, т.к. получается, что мы вычли его два раза вместе с голубым и фиолетовым прямоугольниками. Все эти слагаемые нам известны, т.к. это и есть префиксные суммы, насчитанные изначально! Поэтому и сейчас все будет работать быстро!
Все, что нам осталось сделать это написать цикл, который будет идти по каждой диагонали сверху вниз и проверять каждый квадрат на наличие нулей. Если 0 есть, спускаемся к следующему квадрату претенденту. Если 0 не обнаружен, то расширяем квадрат, увеличивая длину его сторон на один, за счет расширения вниз.
Ну вот и все решение! Префиксные суммы не только прошли все 28 тестов, но и продемонстрировали кратное превосходство по скорости над жадным алгоритмом. Это видно по таблицам ниже, из яндекс контест.
Жадный алгоритм:
| Вердикт | Ресурсы | Баллы |
1 | ok | 45ms | 4.64Mb |
2 | ok | 46ms | 4.51Mb |
3 | ok | 46ms | 4.51Mb |
4 | ok | 0.659s | 4.64Mb |
5 | ok | 0.645s | 4.64Mb |
6 | ok | 0.655s | 4.51Mb |
7 | ok | 0.569s | 4.64Mb |
8 | ok | 0.548s | 4.64Mb |
9 | ok | 58ms | 4.64Mb |
10 | ok | 87ms | 4.64Mb |
11 | time-limit-exceeded | 4.075s | 8.13Mb |
Алгоритм с префиксными суммами:
№ | Вердикт | Ресурсы | Баллы |
1 | ok | 117ms | 33.52Mb |
2 | ok | 116ms | 33.77Mb |
3 | ok | 117ms | 33.52Mb |
4 | ok | 159ms | 36.61Mb |
5 | ok | 159ms | 33.65Mb |
6 | ok | 160ms | 33.65Mb |
7 | ok | 169ms | 37.64Mb |
8 | ok | 163ms | 37.64Mb |
9 | ok | 142ms | 33.39Mb |
10 | ok | 191ms | 39.87Mb |
11 | ok | 301ms | 50.01Mb |
12 | ok | 209ms | 42.41Mb |
13 | ok | 177ms | 39.36Mb |
14 | ok | 146ms | 33.64Mb |
15 | ok | 155ms | 33.52Mb |
16 | ok | 173ms | 39.68Mb |
17 | ok | 288ms | 47.67Mb |
18 | ok | 254ms | 45.99Mb |
19 | ok | 240ms | 44.02Mb |
20 | ok | 185ms | 39.96Mb |
21 | ok | 317ms | 56.85Mb |
22 | ok | 418ms | 57.75Mb |
23 | ok | 408ms | 56.90Mb |
24 | ok | 415ms | 57.51Mb |
25 | ok | 421ms | 57.50Mb |
26 | ok | 406ms | 56.96Mb |
27 | ok | 179ms | 40.73Mb |
28 | ok | 260ms | 51.31Mb |
Всем спасибо, что дочитали до конца. Полные исходники обоих вариантов решения на python выложил сюда. Также предлагаю посмотреть полный видеоролик с решением, который я смонтировал в библиотеке manim. Его можно посмотреть в ТГ, на rutube или в VK Видео (обещаю, что где-то в середине будет довольно забористо :-)
P.S. Чуть не забыл. Если вы увлекаетесь алгоритмами, то заметили, что я не упомянул ничего о вычислительной сложности ни для одного из двух алгоритмов. Это потому что, во первых я сам не очень хорошо это умею делать, а во вторых я прошу вас в комментариях написать свои варианты. Спасибо! )