Задача с небоскрёбом и яйцами — не бином Ньютона?

На самом деле, он самый. Но обо всём по порядку.

Постановка задачи

Осваиваю питон, решаю всякое на Codewars. Сталкиваюсь с известной задачей про небоскрёб и яйца. Разница лишь в том, что исходные данные — не 100 этажей и 2 яйца, а чуть побольше.
Дано: N яиц, M попыток их бросить, бесконечный небоскрёб.

Определить: максимальный этаж, с которого можно бросить яйцо, не разбив. Яйца сферические в вакууме и, если одно из них не разбилось, упав, например, с 99-го этажа, то остальные тоже выдержат падение со всех этажей меньше сотого.

0 <= N, M <= 20000.
Время прогона двух десятков тестов — 12 секунд.

Поиск решения

Нужно написать функцию height(n, m), которая будет возвращать номер этажа для заданных n,m. Так как упоминаться она будет очень часто, а писать каждый раз «height» лень, то везде, кроме кода обозначу её, как f(n, m).

Начнём с нулей. Очевидно, что если нет яиц или попыток их бросить, то определить ничего нельзя и ответом будет ноль. f(0, m) = 0, f(n, 0) = 0.

Допустим, яйцо одно, а попыток 10. Можно рискнуть всем и бросить его сразу с сотого этажа, но в случае неудачи, определить больше ничего не удастся, так что логичнее всё же начинать с первого этажа и подниматься вверх на один этаж после каждого броска, пока не кончатся либо попытки, либо яйцо. Максимум, куда получится добраться, если яйцо не подведёт — этаж номер 10. f(1, m) = m

Возьмём второе яйцо, попыток опять 10. Теперь-то можно рискнуть с сотого? Разобьётся — останется ещё одно и 9 попыток, хотя бы 9 этажей получится пройти. Так может и рисковать нужно не с сотого, а с десятого? Логично. Тогда, в случае успеха, останется 2 яйца и 9 попыток. По аналогии, теперь нужно подняться ещё на 9 этажей. При серии успехов — ещё на 8, 7, 6, 5, 4, 3, 2 и 1. Итого, мы на 45-ом этаже с двумя целыми яйцами и без попыток. Ответом видится сумма первых M членов арифметической прогрессии с первым членом 1 и шагом 1. f(2, m) = (m*m + m) / 2. А ещё видно, что на каждом шаге как бы происходил вызов функции f(1, m), но это пока не точно.

Продолжим с тремя яйцами и десятью попытками. В случае неудачного первого броска, прикрыты снизу будут этажи, определяемые 2 яйцами и 9 попытками, а это значит, что первый бросок надо совершать с этажа f(2, 9) + 1. Тогда, в случае успеха, имеем 3 яйца и 9 попыток. И для второй попытки нужно подняться ещё на f(2,8) + 1 этажей. И так далее, пока на руках не останется 3 яйца и 3 попытки. И тут самое время отвлечься на рассмотрение случаев с N=M, когда яиц столько же, сколько и попыток.

Вернёмся к f(3, 10). По сути, на каждом шаге всё сводится к сумме f(2, m-1) — количеству этажей, которые можно определить в случае неудачи, единицы и f(3, m-1) — количеству этажей, которые можно определить в случае успеха. И становится понятно, что от увеличения количества яиц и попыток вряд ли что-то изменится. f(n, m) = f(n — 1, m — 1) + 1 + f(n, m — 1). И это универсальная формула, можно воплощать в коде.

from functools import lru_cache
@lru_cache()
def height(n,m):
    if n==0 or m==0: return 0
    elif n==1: return m
    elif n==2: return (m**2+m)/2
    elif n>=m: return 2**n-1
    else: return height(n-1,m-1)+1+height(n,m-1)

Разумеется, предварительно я наступил на грабли немемоизирования рекурсивных функций и узнал, что f(10, 40) выполняется почти 40 секунд с количеством вызов самой себя — 97806983. Но и мемоизация спасает только на начальных интервалах. Если f(200,400) выполняется за 0.8 секунды, то f(200, 500) уже за 31 секунду. Забавно, что при замере времени выполнения с помощью %timeit, то результат получается куда меньше реального. Очевидно, что первый прогон функции занимает большую часть времени, а остальные просто пользуются результатами его мемоизации. Ложь, наглая ложь и статистика.

Рекурсия не нужна, ищем дальше

Итак, в тестах фигурирует, например, f(9477, 10000), а у меня жалкая f(200, 500) уже не укладывается в нужное время. Значит есть другое решение, без рекурсии, продолжим его поиски. Дополнил код подсчетом вызовов функции с определёнными параметрами, чтобы посмотреть на что она в итоге раскладывается. Для 10 попыток получились такие результаты:

f(3,10) = 7+ 1*f(2,9)+ 1*f(2,8)+ 1*f(2,7)+ 1*f(2,6)+ 1*f(2,5)+ 1*f(2,4)+ 1*f(2,3)+ 1*f(3,3)
f(4,10) = 27+ 1*f(2,8)+ 2*f(2,7)+ 3*f(2,6)+ 4*f(2,5)+ 5*f(2,4)+ 6*f(2,3)+ 6*f(3,3)+ 1*f(4,4)
f(5,10) = 55+ 1*f(2,7)+ 3*f(2,6)+ 6*f(2,5)+ 10*f(2,4)+ 15*f(2,3)+ 15*f(3,3)+ 5*f(4,4)+ 1*f(5,5)
f(6,10) = 69+ 1*f(2,6)+ 4*f(2,5)+ 10*f(2,4)+ 20*f(2,3)+ 20*f(3,3)+ 10*f(4,4)+ 4*f(5,5)+ 1*f(6,6)
f(7,10) = 55+ 1*f(2,5)+ 5*f(2,4)+ 15*f(2,3)+ 15*f(3,3)+ 10*f(4,4)+ 6*f(5,5)+ 3*f(6,6)+ 1*f(7,7)
f(8,10) = 27+ 1*f(2,4)+ 6*f(2,3)+ 6*f(3,3)+ 5*f(4,4)+ 4*f(5,5)+ 3*f(6,6)+ 2*f(7,7)+ 1*f(8,8)
f(9,10) = 7+ 1*f(2,3)+ 1*f(3,3)+ 1*f(4,4)+ 1*f(5,5)+ 1*f(6,6)+ 1*f(7,7)+ 1*f(8,8)+ 1*f(9,9)

Какая-то закономерность просматривается:

Коэффициенты эти теоретически рассчитываются. Каждый голубой — сумма верхнего и левого. А фиолетовые — те же самые голубые, только в обратном порядке. Рассчитать можно, но это же опять рекурсия, а в ней я разочаровался. Скорее всего, многие (жаль, что не я) уже узнали эти числа, но я пока сохраню интригу, следуя своему пути решения. Решил на них плюнуть и зайти с другой стороны.

Открыл эксель, соорудил табличку с результатами работы функции и стал высматривать закономерности. С3=ЕСЛИ(C$2>$B3;2^$B3-1;C2+B2+1), где $2 — строка с количеством яиц (1-13), $B — столбец с количеством попыток (1-20), C3 — ячейка на пересечении двух яиц и одной попытки.

Серая диагональ это N=M и тут наглядно видно, что справа от неё (при N > M) ничего не меняется. Видно — да иначе и быть не может, ведь это все результаты работы формулы, в которой так и задано- что каждая ячейка равна сумме верхней, верхней слева и единицы. Но какой-то универсальной формулы, куда можно подставить N и M и получить номер этажа, найти не удалось. Спойлер: её и не существует. Но зато так просто создание этой таблицы в экселе выглядит, может быть можно сгенерировать такую же питоном и таскать из неё ответы?

NumPy, вам нет

Вспоминаю, что существует NumPy, который как раз предназначен для работы с многомерными массивами, почему бы не попробовать его? Для начала нам понадобится одномерный массив из нулей размером в N+1 И одномерный массив из единиц размером в N. Берём первый массив от нулевого до предпоследнего элемента, складываем поэлементно с первым же массивом от первого элемента до последнего и с массивом единиц. К получившемуся массиву в начало добавляем ноль. Повторять M раз. Элемент номер N получившегося массива и будет ответом. Первые 3 шага выглядят так:

NumPy работает так быстро, что сохранять всю таблицу я не стал — каждый раз считал нужную строку заново. Одно но — результат работы на больших числах был неправильный. Старшие разряды вроде те, а младшие — нет. Так выглядят накопившиеся от многократного сложения погрешности арифметики чисел с плавающей точкой. Не беда — можно же изменить тип массива на int. Нет, беда — оказалось, что NumPy ради быстродействия работает только со своими типами данных, а его int, в отличие от питоновского int, может быть не больше 2^64-1, после чего молча переполняется и продолжается с -2^64. А у меня вообще-то ожидаются числа под три тысячи знаков. Зато работает очень быстро, f(9477, 10000) выполняется 233 мс, просто на выходе ерунда какая-то получается. Код даже приводить не буду, раз такое дело. Попробую сделать то же самое чистым питоном.

Итерировал, итерировал, да не выитерировал

def height(n, m):
    arr = [0]*(n+1)
    while m > 0:
        arr = [0] + list(map(lambda x,y: x+y+1, arr[:-1], arr[1:]))
        m-=1
    return arr[n]

44 секунды на расчет f(9477, 10000), многовато. Но зато абсолютно точно. Что можно оптимизировать? Первое — нет нужды считать всё, что правее диагонали M,M. Второе — ни к чему считать последний массив целиком, ради одной-то ячейки. Для этого сгодятся две последние две ячейки предыдущего. Для вычисления f(10, 20) будет достаточно только вот этих серых ячеек:

И так это выглядит в коде:

def height(n, m):
    arr = [0, 1, 1]
    i = 1
    while i < n and i < m-n: # треугольник слева от m,m
        arr = [0] + list(map(lambda x,y: x+y+1, arr[:-1], arr[1:]))
        arr += [arr[-1]]
        i+=1    
    arr.pop(-1)
    while i < n or i < m-n: # прямоугольник или трапеция до начала следующего треугольника
        arr = list(map(lambda x,y: x+y+1, arr[:-1], arr[1:]))
        arr = arr + [arr[-1]+1] if n > len(arr) else [0] + arr
        i+=1
    while i < m: # треугольник, сходящийся в одну ячейку - ответ
        arr = list(map(lambda x,y: x+y+1, arr[:-1], arr[1:]))
        i+=1
    return arr[0]

И что вы думаете? f(9477, 10000) за 2 секунды! Но эти входные данные слишком хороши, длина массива на любом этапе будет не больше 533 элементов (10000-9477). Проверим на f(5477, 10000) — 11 секунд. Тоже хорошо, но только по сравнению с 44 секундами — двадцать тестов с таким временем не пройдёшь.

Всё не то. Но раз есть задача, значит есть и решение, поиски продолжаются. Я снова стал смотреть в экселевскую таблицу. Ячейка слева от (m, m) всегда на единицу меньше. А ячейка слева от неё уже нет, в каждой строке разность становится больше. Ячейка снизу от (m, m) всегда вдвое больше. А ячейка снизу от неё — уже не вдвое, а чуть меньше, но для каждого столбца по разному, чем дальше, тем больше. А ещё числа в одной строке сначала растут быстро, а после середины медленно. Дай-ка я построю таблицу разностей между соседними ячейками, может там какая закономерность появится?

Теплее

Ба, знакомые числа! То есть сумма N этих числе в строке номер M это и есть ответ? Правда, считать их — примерно то же самое, что я уже делал, вряд ли это сильно ускорит работу. Но надо попробовать:

def height(n, m):
    arr = [1,1]
    while m > 1:
        arr = [1] + list(map(lambda x,y: x+y, arr[1:], arr[:-1])) + [1]
        m-=1
    return sum(arr[1:n+1])

def height(n, m):
    arr = [1,1]
    while m > 2 and len(arr) < n+2: # половина треугольника или трапеция, если n < половины
        arr = [1] + list(map(lambda x,y: x+y, arr[1:], arr[:-1]))
        arr = [1] + list(map(lambda x,y: x+y, arr[1:], arr[:-1]))
        arr += [arr[-1]]
        m-=2
    while m > 1: 
        arr = [1] + list(map(lambda x,y: x+y, arr[1:], arr[:-1]))
        m-=1
    if len(arr) < n+1: arr += arr[::-1][1:] # если n во второй половине, отражаем строку
    return sum(arr[1:n+1])

Не сказать, что более оптимальное решение. На каких-то данных работает быстрее, на каких-то медленнее. Надо отправляться глубже. Что это за треугольник с числами, что появился в решении дважды? Стыдно признаться, но высшую математику, где треугольник наверняка фигурировал, я благополучно позабыл, так что пришлось гуглить.

Бинго!

Треугольник Паскаля, так он официально называется. Бесконечная таблица биномиальных коэффициентов. Так что ответом на задачу с N яйцами и M бросками является сумма первых N коэффициентов в разложении бинома Ньютона M-ой степени, кроме нулевого.

Произвольный биномиальный коэффициент может быть вычислен через факториалы номера строки и номера коэффициента в строке: bk = m!/(n!*(m-n!)). Но самое приятное — можно последовательно рассчитать числа в строке, зная её номер и нулевой коэффициент (всегда единица): bk[n] = bk[n-1] * (m — n + 1) / n. На каждом шаге числитель уменьшается на единицу, а знаменатель увеличивается. И лаконичное финальное решение выглядит так:

def height(n, m):
    h, bk = 0, 1 # высота и нулевой биномиальный коэффициент
    for i in range(1, n + 1): 
        bk = bk * m // i 
        h += t
        m-=1
    return h

33 мс. на расчёт f(9477, 10000)! Это решение тоже можно оптимизировать, хотя в заданных диапазонах и оно отлично работает. Если n лежит во второй половине треугольника, то можно инвертировать его в m-n, посчитать сумму первых n коэффициентов и отнять её от 2^m-2. Если n близко к середине и m нечётное, то расчеты тоже можно сократить: сумма первой половины строки будет равна 2^(m-1)-1, последний коэффициент в первой половине можно рассчитать через факториалы, его номер это (m-1)/2, а дальше либо продолжаем прибавлять коэффициенты, если n в правой половине треугольника, либо отнимать, если в левой. Если m чётное, то половину строки уже не посчитаешь, но можно найти сумму первых m/2+1 коэффициентов, посчитав через факториалы средний и прибавив половину его к 2^(m-1)-1. На входных данных в районе 10^6 это очень заметно уменьшает время выполнения.

Уже после успешного решения, стал искать чьи-нибудь ещё изыскания по этому вопросу, но нашёл только то самое, с собеседований, с всего лишь двумя яйцами, а это не спортивно. Интернет будет неполон без моего решения, решил я, и вот оно.

Автор: nullemotion

Источник