Как стать автором
Обновить

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

Уровень сложностиПростой
Время на прочтение5 мин
Количество просмотров4.3K

Расскажу о том, как решал одну из наиболее интересных задач в разминке Яндекс Алгоритмы 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. Чуть не забыл. Если вы увлекаетесь алгоритмами, то заметили, что я не упомянул ничего о вычислительной сложности ни для одного из двух алгоритмов. Это потому что, во первых я сам не очень хорошо это умею делать, а во вторых я прошу вас в комментариях написать свои варианты. Спасибо! )

Теги:
Хабы:
Всего голосов 7: ↑6 и ↓1+6
Комментарии30

Публикации

Работа

Data Scientist
49 вакансий

Ближайшие события