求解差分方程/函数方程优化DP：N个鸡蛋掉落问题

摘要: 可以直接作为差分方程/函数方程求解的状态转移方程

【对算法，数学，计算机感兴趣的同学，欢迎关注我哈，阅读更多原创文章】
我的网站：潮汐朝夕的生活实验室
我的公众号：算法题刷刷
我的知乎：潮汐朝夕
我的github：FennelDumplings
我的leetcode：FennelDumplings

---

对于某些简单的状态转移方程，可以直接作为差分方程/函数方程求解出来。求通项公式很可能会比直接推导状态时间复杂度低。这种方法需要一些灵感。可以参考 推公式优化DP。

鸡蛋掉落问题的值域二分算法中，判定的时候用的是动态规划，其状态转移方程是可以作为差分方程求解的，两个鸡蛋的情况比较简单，参考文章 两枚鸡蛋的鸡蛋掉落：单峰性、决策单调性、差分方程。

上面两个鸡蛋的情况发现差分方程的解比较容易，N个鸡蛋情况的状态转移方程，作为差分方程求解就不容易了，需要对常见的函数方程的解有所了解，本文来看这个问题。

887. 鸡蛋掉落

给你 k 枚相同的鸡蛋，并可以使用一栋从第 1 层到第 n 层共有 n 层楼的建筑。

已知存在楼层 f ，满足 0 <= f <= n ，任何从 高于 f 的楼层落下的鸡蛋都会碎，从 f 楼层或比它低的楼层落下的鸡蛋都不会破。

每次操作，你可以取一枚没有碎的鸡蛋并把它从任一楼层 x 扔下（满足 1 <= x <= n）。如果鸡蛋碎了，你就不能再次使用它。如果某枚鸡蛋扔下后没有摔碎，则可以在之后的操作中 重复使用 这枚鸡蛋。

请你计算并返回要确定 f 确切的值 的 最小操作次数 是多少？

提示：

1 <= k <= 100
1 <= n <= 1e4

示例 1：
输入：k = 1, n = 2
输出：2
解释：
鸡蛋从 1 楼掉落。如果它碎了，肯定能得出 f = 0 。
否则，鸡蛋从 2 楼掉落。如果它碎了，肯定能得出 f = 1 。
如果它没碎，那么肯定能得出 f = 2 。
因此，在最坏的情况下我们需要移动 2 次以确定 f 是多少。

示例 2：
输入：k = 2, n = 6
输出：3

示例 3：
输入：k = 3, n = 14
输出：4

算法：值域二分 + 动态规划

我们要求的是，给定鸡蛋个数 $N$ 和楼层 $H$，测完所有未知范围所需的最小操作次数。

如果考虑给定操作次数 $m$，给定鸡蛋个数 $N$，最多可以测多少楼层，记其为 $f(m, N)$，通过问题背景可以知道 $f(m, N)$ 随 $m$ 单调增加，即操作步数 $m$ 越多，则可以测的楼层 $f(m, N)$ 越多。

因此可以用值域二分算法。初始时 $[left, right] = [1, H]$，如果 $f(mid, N) < H$，则答案一定大于 $mid$，否则答案一定小于等于 $mid$。

下面考虑 $f(m, N)$ 如何求。定义 $dp[i][k]$ 表示 $i$ 步操作，还剩 $k$ 个鸡蛋，可以测的最多楼层。那么 $f(m, N) = dp[m][N]$。初始情况是 $dp[i][0] = 0, dp[0][k] = 0$，即步数为零或鸡蛋为零时，能测出的楼层均为零。

考察还剩 $i$ 步操作，此时有 $k$ 个鸡蛋，假设测的楼层为 $x$：

• 如果没碎，那么 $x$ 及其下面的层均已知，鸡蛋依然是 $k$ 个，$x$ 上面可以再测 $dp[i - 1][k]$ 层。如果 $x$ 下面有 $a$ 层，一共测出 $1 + dp[i-1][k] + a$ 层。
• 如果碎了，那么 $x$ 及其上面的层均已知，还剩 $k-1$ 个鸡蛋，$x$ 下面可以再测 $dp[i - 1][k - 1]$ 层。如果 $x$ 上面有 $b$ 层，一共测出 $1 + dp[i - 1][k - 1] + b$ 层。

给定了 $a, b$ 之后，不论鸡蛋是否碎，$x$ 下面总能测出的层数为 $\min(dp[i-1][k-1], a)$、$x$ 上面总能测出的层数为 $\min(b, dp[i - 1][k])$，汇总起来，有：

$$dp[i][k] = \max\limits_{a, b}(1 + \min(dp[i-1][k-1], a) + \min(b, dp[i - 1][k]))$$

$g(a) = \min(dp[i-1][k-1], a), h(b) = \min(b, dp[i - 1][k])$ 的图像如下：

可以直观看出 $dp[i][k] = 1 + dp[i - 1][k-1] + dp[i - 1][k]$。

在值域二分的过程中如果按照上述动态规划的流程求 $f(m, N)$，则最坏需要 $O(HN)$，这样值域二分的时间复杂度为 $O(HN\log H)$。

---

但是这里有个性质，到在算 $f(m, N)$ 的过程中，实际上 $f(1, N), f(2, N), \cdots, f(m, N)$ 都算出来了。假设当前二分中点为 $mid$，则计算 $f(mid, N)$ 的过程顺便就算了 $f(1, N),f(2, N),\cdots,f(mid-1, N)$，过程中算到哪个 $f(i, N) \geq H$，直接就可以返回 $i$ 而不用等到 $mid$ 算完；如果算完了 $f(mid, N)$ 依然小于 $H$，那么就需要继续算 $f(mid + 1, N), f(mid + 2, N),\cdots$。

因此，这里可以直接按照 $f(1, N), f(2, N), \cdots$ 来算，算到某个 $f(i, N) \geq H$ 时返回即可。

假设 $i$ 推导到 $c$ 时 $dp[c][N] \geq H$，则时间复杂度为 $O(Nc)$。

代码 (C++)

class Solution {
public:
    int superEggDrop(int N, int H) {
        vector<vector<int> > dp(H + 1, vector<int>(N + 1, 0));
        for (int i = 1; i <= H; ++i) 
        {
            dp[i][0] = 0; 
            for (int k = 1; k <= N; ++k) 
            {
                dp[i][k] = dp[i - 1][k - 1] + dp[i - 1][k] + 1;
                if (dp[i][k] >= H) 
                    return i;
            }
        }
        return H;
    }
};

代码 (Python)

class Solution:
    def superEggDrop(self, N: int, H: int) -> int:
        dp = [[0 for _ in range(N + 1)] for _ in range(H + 1)]
        for i in range(1, H + 1):
            for k in range(1, N + 1):
                dp[i][k] = dp[i - 1][k - 1] + dp[i - 1][k] + 1;
                if dp[i][k] >= H: 
                    return i
        return N

二元差分方程优化 DP

回顾状态转移方程：

$$\begin{align} &dp[i][k] = 1 + dp[i - 1][k-1] + dp[i - 1][k] \\ &dp[i][0] = 0 \\ &dp[0][k] = 0 \\ \end{align}$$

可以将其作为差分方程直接求解，如下：

$$\begin{align} &f(m, N) = 1 + f(m - 1, N-1) + f(m - 1, N) \\ &f(m, 1) = m \\ &f(1, N) = 1 \\ \end{align}$$

构造 $f(m, N)$ 关于 $N$ 的差分 $g(m, N) = f(m, N) - f(m, N - 1)$：

$$f(m, N-1) = 1 + f(m - 1, N-2) + f(m - 1, N-1)$$

将 $f(m, N), f(m, N - 1)$ 分别代入 $g(m, N)$，得到：

$$g(m, N) = f(m-1, N) - f(m-1, N-2)$$

另一方面：

$$g(m-1, N) = f(m-1, N) - f(m-1, N - 1)$$ $$g(m-1, N-1) = f(m-1, N-1) - f(m-1, N - 2)$$

因此得到关于差分 $g(m, N)$ 的方程：

$$g(m, N) = g(m - 1, N) + g(m - 1, N - 1)$$

这是一个有组合数学背景的差分方程，正是组合数 $\binom{m}{N}$ 的递推公式。因此 $g(m, N) = \binom{m}{N}$。

由于 $g(m, N)$ 是 $f(m, N)$ 关于 $N$ 的差分，因此 $f(m, N)$ 很容易通过 $g(m, N)$ 求和得到：

$$\begin{align} f(m, N) &= f(m, 1) + g(m, 2) + g(m, 3) + \cdots + g(m, N) \\ &= m + \sum\limits_{k=2}^{N}g(m, k) \\ &= m + \sum\limits_{k=2}^{N}\binom{m}{k} \\ &= \sum\limits_{k=1}^{N}\binom{m}{k} \\ \end{align}$$

由于 $m < N$ 时 $\binom{m}{N} = 0$，$f(m, N)$ 可以简化为：

$$f(m, N) = \sum\limits_{k=1}^{\min(m, N)}\binom{m}{k} \\$$

在算上述求和的时候，当中间某一步和式已经大于等于 $H$ 的时候，就可以直接返回了，因为二分过程仅需要判断 $f(m, N)$ 是否小于 $H$。

代码(Python)

可以先把 $m \in [1, H]$ 范围内的组合数 $\binom{m}{1},\cdots,\binom{m}{N}$ 先预处理出来，时间复杂度为 $O(NH)$，然后在二分的时候，以 $O(N)$ 做个求和即可判定，时间复杂度 $O(N\log H)$，总时间复杂度 $O(NH)$。

class Solution:
    def superEggDrop(self, N: int, H: int) -> int:
        C = [[0 for _ in range(N + 1)] for _ in range(H + 1)]
        for m in range(H + 1):
            C[m][0] = 1
        for m in range(1, H + 1):
            for k in range(1, min(H, N) + 1):
                C[m][k] = C[m - 1][k - 1] + C[m - 1][k]

        def f(m, N):
            ans = 0
            for k in range(1, min(m, N) + 1):
                ans += C[m][k]
            return ans

        left = 1
        right = H
        while left < right:
            mid = (left + right) // 2
            if f(mid, N) < H:
                left = mid + 1
            else:
                right = mid
        return left

代码 (Python)

如果不预处理而是在算到 $\binom{m}{k}$ 时，直接用阶乘的公式算（math.comb(m, k) 的算法）：

$$\binom{m}{k} = \frac{m!}{k!(m-k)!} = \frac{m}{k}\cdot\frac{m-1}{k-1}\cdot\cdot\cdot\frac{m-k+1}{1}$$

判定时需要对 $k=1,\cdots,N$ 求上式并求和，时间复杂度为 $O(N^{2})$，总时间复杂度为 $O(N^{2}\log H)$。

class Solution:
    def superEggDrop(self, N: int, H: int) -> int:
        def f(m: int, N: int) -> int:
            ans = 0
            for k in range(1, min(m, N) + 1):
                ans += comb(m, k)
                if ans >= H:
                    break
            return ans

        left = 1
        right = H
        while left < right:
            mid = (left + right) // 2
            if f(mid, N) < H:
                left = mid + 1
            else:
                right = mid
        return left

组合数的循环求和的优化

回顾公式：

$$f(m, N) = \sum\limits_{k=1}^{\min(m, N)}\binom{m}{k} \\$$

注意到 $\binom{m}{k} = \frac{m-k+1}{k}\binom{m}{k-1}$，于是在求和过程中，算到 $k$ 这一项的时候 $\binom{m}{k}$ 不用单独算，直接一步乘法 $\frac{m-k+1}{k}\binom{m}{k-1}$ 即可，而 $\binom{m}{k-1}$ 在求和算 $k-1$ 这一项的时候已经算过，因此求和的项可以在 $O(1)$ 算出，真个求和的时间复杂度就是 $O(N)$。

于是整体时间复杂度为 $O(N\log H)$。

代码 (Python)

class Solution:
    def superEggDrop(self, N: int, H: int) -> int:
        def f(m: int, N: int) -> int:
            ans = 0
            # 上一轮求和项 C(m, k-1) 的值
            # k = 1 时，C(m, k-1) = 1
            r = 1 
            for k in range(1, min(m, N) + 1):
                r *= (m - k + 1)
                r //= k
                ans += r
                if ans >= H:
                    break
            return ans

        left = 1
        right = H
        while left < right:
            mid = (left + right) // 2
            if f(mid, N) < H:
                left = mid + 1
            else:
                right = mid
        return left

代码 (C++)

class Solution {
public:
    int superEggDrop(int N, int H) 
    {
        int left = 1, right = H;
        while (left < right) 
        {
            int mid = (left + right) / 2;
            if (f(mid, N, H) < H)
                left = mid + 1;
            else
                right = mid;
        }
        return left;
    }

private:
    int f(int m, int N, const int H) 
    {
        int ans = 0;
        // 上一轮求和项 C(m, k-1) 的值
        // k = 1 时，C(m, k-1) = 1
        int r = 1; 
        for (int k = 1; k <= min(m, N); ++k) 
        {
            r *= m - k + 1;
            r /= k;
            ans += r;
            if(ans >= H)
                break;
        }
        return ans;
    }
};

---