[数解算法] 扔鸡蛋问题

扔鸡蛋问题是一道经典的算法题，LeetCode链接。

一般做法是动态规划，时间复杂度 $O(KF^2)$ ，对应K个蛋F层楼。
优化后可以达到 $O(KF)$ 。

该题目下有很多算法详解，包括官方解答。

但是今天我们不聊动态规划，从数学的角度解决这个问题。

初步分析

有 $k$ 个鸡蛋， $f$ 层楼，求最小次数 $n(k,f)$ 。

易知， $n(1,f) = f$ ，即一个鸡蛋只能从一楼开始逐层尝试。

然而，当 $k=2$ 的时候，问题就难以通过简单的心算来得到答案。但是容易观察到：

$n(2,2) = 2$
$n(2,3) = 2$

即，2层楼和3层楼都需要扔2次。这一现象也很直观，不考虑蛋碎的情况，2次二分查找最多可以遍历3项。

同时，我们发现函数 $n(2,f)$ 不是严格单调的，而是分段阶梯式，这非常不利于我们求解。

所以，我们不妨转化为问题 $f(k,n)$ ，即给定鸡蛋和次数，最多可以遍历的楼层数。如果该问题得解，我们很容易求解原问题。

最多楼层

现在我们不需要标记楼层号了，而只需要知道自己遍历了多少层。

若当前有 $k$ 蛋， $n$ 次机会，则在一次丢蛋后

剩余机会必然为 $n-1$
蛋可能为 $k$ 或者 $k-1$ ，对应向上和向下继续遍历
同时，当前楼层已经被遍历

所以我们有了下式

\begin{align} f(k,n) &= 1+ f(k-1, n-1) + f(k, n-1) \end{align}

又由于1个蛋的情况我们已经得解，即

\begin{align} f(1,n) &= n \end{align}

把 $(2)$ 带入 $(1)$ 中，可得

\begin{aligned} f(2,n) &= 1+ f(1,n-1) + f(2, n-1) \\ &= 1 + n - 1 + f(2, n-1) \\ &= n + f(2,n-1) \end{aligned}

这不就变成了等差数列求和，所以我们有

\begin{align} f(2,n) &= \frac 12 n^2 + \frac 12 n \end{align}

如果我们把 $(3)$ 也带入 $(1)$ 中，即可得 $f(3,n)$ 的通项。

虽然我们可以求解，但这离我们的通解还有一些距离。

求解

虽然我们无法直接求解通项，但是易知

(我们把 $f$ 看做 $n$ 的多项式，而 $k$ 为常数)

$f(k,n)$ 为 $k$ 次多项式
$f(k,n)$ 过点 $(i,2^i-1)), 0<=i<=k$

这么一来，我们就有了 $k+1$ 个方程求解 $k$ 次方程。

记该多项式 $i$ 次幂的系数为 $a_{k,i}$ ，即有：

\begin{align} A_k & = \begin{bmatrix} a_{k,0} & a_{k,1} & \cdots & a_{k,k} \end{bmatrix} ^ T \\ R_k & = \begin{bmatrix} 0 & 1 & 3 & \cdots & 2^k - 1 \end{bmatrix} ^ T \\ M_k & = \begin{pmatrix} 0^0 & 0^1 & \cdots & 0^k \\ 1^0 & 1^1 & \cdots & 1^k \\ \vdots & \vdots & \ddots & \vdots \\ k^0 & k^1 & \cdots & k^k \end{pmatrix} \\ M_k \cdot A_k &= R_k \\ A_k &= M_k^{-1} R_k \end{align}

现在，我们就得到了解，对于 $k$ 个蛋， $n$ 次机会，最多遍历楼层数为：

\begin{align} f(k,n) = \begin{bmatrix} n^0 & n^1 & \cdots & n^k \end{bmatrix} \cdot A_k \end{align}

该解法的时间复杂度为 $O(k^3)$ ，即矩阵求逆。

但是考虑到，如果 $n<k或者f<2^k$ 时，可以直接求解，故

O(k^3) < O(k \cdot ln^2f)

答题

至此，我们已经得到了题解，接下来我们把它转化为代码提交给LeetCode。

以下代码可以直接提交给LeetCode通过。

import numpy as np
import numpy.linalg as nl

class Solution:
    def superEggDrop(self, egg: int, floor: int) -> int:
        if 2 ** egg > floor:
            return int(np.ceil(np.log2(floor + 1)))
        M = np.array([[i ** j for j in range(egg + 1)] for i in range(egg + 1)])
        R = np.array([[2 ** i - 1] for i in range(egg + 1)])
        A = np.matmul(nl.inv(M), R)
        A[0] = -floor
        roots = np.roots(np.flip(A.T[0]))
        root = [e.real for e in roots if np.isreal(e) and e.real > 0][0].real
        return int(np.ceil(np.round(root, 6)))

其中M, R, A对应前式相同符号。唯一需要注意的是，python求根有精度误差，所以在结果返回前进行了一次6位精度的round。

我的解答的运行时间为76ms，对比官方答案需要184ms，相差三倍。

后记

在写完本文后，我查看了LeetCode上别人的解法。受到这个答案的启发，才发现自己组合数学确实不好。

式 $(1)$ 其实有非常简洁的形式。

我们已知组合恒等式

\begin{align} {n \choose k} = {n-1 \choose k-1} + {n-1 \choose k} \end{align}

这与式 $(1)$ 非常相似。所以我们需要从中构造出一个多余的1。

将式 $(10)$ 从 $k=k$ 向 $k=1$ 展开，即

\begin{align} {n \choose k} &= {n-1 \choose k-1} + {n-1 \choose k} \\ {n \choose k-1} &= {n-1 \choose k-2} + {n-1 \choose k-1} \\ &\vdots \\ {n \choose 2} &= {n-1 \choose 1} + {n-1 \choose 2} \\ {n \choose 1} &= {n-1 \choose 0} + {n-1 \choose 1} \end{align}

我们将 $k$ 个式子相加，右边最后一个式子的第一项 ${n-1 \choose 0}$ 正好是我们需要的 $1$ ，而去掉这一项后，此列正好剩下 $k-1$ ，正好满足我们的形式，即

\begin{align} \sum_{i=1}^k {n \choose i} &= 1+ \sum_{i=1}^{k-1} {n-1 \choose i} + \sum_{i=1}^k {n-1 \choose i} \end{align}

所以，我们得到了式 $(1)$ 的另一个形式

\begin{align} f(k,n) &= \sum_{i=1}^k {n \choose i} \end{align}

import numpy as np
from scipy.special import comb

class Solution:
    def superEggDrop(self, egg: int, floor: int) -> int:
        if 2 ** egg > floor:
            return int(np.ceil(np.log2(floor + 1)))
        start = int(np.floor(np.log2(floor)))
        for n in range(start, floor + 1):
            if sum([comb(n, e + 1) for e in range(0,egg)]) >= floor:
                return int(n)

然而，这种解法并不会比上面的多项式解法快，因为计算一次值得复杂度就已经是 $O(n^3)$ ，而我们难以估计精确值，需要在较大的区间内搜索。

在LeetCode中运行，耗时浮动在100ms~150ms。耗时比前一种解法多50%。