Префиксные суммы. Решение задачи из тренировок Яндекса по алгоритмам

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

И так задача:

Кролики очень любопытны. Они любят изучать геометрию, бегая по грядкам. Наш кролик как раз такой. Сегодня он решил изучить новую фигуру — квадрат.

Кролик бегает по грядке — клеточному полю N × M клеток. В некоторых из них посеяны морковки, в некоторых нет.

Помогите кролику найти сторону квадрата наибольшей площади, заполненного морковками полностью.

Формат ввода

В первой строке даны два натуральных числа N и M ( 1 ≤ N, M ≤ 1000). Далее в N строках расположено по M чисел, разделенных пробелами (число равно 0, если в клетке нет морковки или 1, если есть).

Формат вывода

Выведите одно число — сторону наибольшего квадрата, заполненного морковками.

Пример 1

Ввод

4 5

Ответ: 2

Пример 2

Ввод

5 5

Ответ: 5

Пример 3

Ввод

4 4

Ответ: 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 тестов, но и продемонстрировали кратное превосходство по скорости над жадным алгоритмом. Это видно по таблицам ниже, из яндекс контест.

Жадный алгоритм:

Алгоритм с префиксными суммами:

Всем спасибо, что дочитали до конца. Полные исходники обоих вариантов решения на python выложил сюда. Также предлагаю посмотреть полный видеоролик с решением, который я смонтировал в библиотеке manim. Его можно посмотреть в ТГ, на rutube или в VK Видео (обещаю, что где-то в середине будет довольно забористо 🙂

P.S. Чуть не забыл. Если вы увлекаетесь алгоритмами, то заметили, что я не упомянул ничего о вычислительной сложности ни для одного из двух алгоритмов. Это потому что, во первых я сам не очень хорошо это умею делать, а во вторых я прошу вас в комментариях написать свои варианты. Спасибо! )