Leetcode 1420：生成数组（超详细的解法！！！）

给你三个整数 n、m 和 k 。下图描述的算法用于找出正整数数组中最大的元素。

请你生成一个具有下述属性的数组 arr ：

• arr 中有 n 个整数。
• 1 <= arr[i] <= m 其中 (0 <= i < n) 。
• 将上面提到的算法应用于 arr ，search_cost 的值等于 k 。

返回上述条件下生成数组 arr 的 方法数 ，由于答案可能会很大，所以 必须 对 10^9 + 7 取余。

示例 1：

输入：n = 2, m = 3, k = 1
输出：6
解释：可能的数组分别为 [1, 1], [2, 1], [2, 2], [3, 1], [3, 2] [3, 3]

示例 2：

输入：n = 5, m = 2, k = 3
输出：0
解释：没有数组可以满足上述条件

示例 3：

输入：n = 9, m = 1, k = 1
输出：1
解释：可能的数组只有 [1, 1, 1, 1, 1, 1, 1, 1, 1]

示例 4：

输入：n = 50, m = 100, k = 25
输出：34549172
解释：不要忘了对 1000000007 取余

示例 5：

输入：n = 37, m = 17, k = 7
输出：418930126

提示：

• 1 <= n <= 50
• 1 <= m <= 100
• 0 <= k <= n

解题思路

比较容易想到递归记忆化的写法，定义函数$f(a,b,c)$表示n=a、k=b、m=c的时候方案数，那么

• $f(a,b,c)={\sum }_{d=c+1}^{m}f(a-1,b-1,d)+f(a-1,b,c)\ast c$

也就是当$d>c$的时候需要search_cost-1。然后就是边界条件，如果$a=0\&b=0$，那么有解。

from functools import lru_cache
class Solution:
    def numOfArrays(self, n: int, m: int, k: int) -> int:
        mod = 10**9 + 7

        @lru_cache(None)
        def dfs(a, b, c):
            if a == 0: 
                return 1 if b == 0 else 0
            
            res = dfs(a - 1, b, c) * c % mod
            for d in range(c + 1, m + 1):
                res = (res + dfs(a - 1, b - 1, d)) % mod
            return res
    
        return dfs(n, k, 0)

上面这个代码想要转化成动态规划的形式不太容易。

class Solution:
    def numOfArrays(self, n: int, m: int, k: int) -> int:
        mod = 10**9 + 7

        dp = [[[0] * (m + 1) for _ in range(k + 1)] for _ in range(n + 1)]
        for c in range(1, m + 1):
            dp[1][1][c] = 1
        for a in range(2, n + 1):
            for b in range(1, k + 1):
                for c in range(1, m + 1):
                    dp[a][b][c] = dp[a - 1][b][c] * c % mod
                    for d in range(c + 1, m + 1):
                        dp[a][b][c] = (dp[a][b][c] + dp[a - 1][b - 1][d]) % mod
        res = 0
        for c in range(1, m + 1):
            res = (res + dp[n][k][c]) % mod
        return res

这个比较难理解的地方就是最后这个累加的过程，难道直接返回dp[-1][-1][-1]不行吗？这实际上是函数定义的问题，我们定义的函数$f(a,b,c)$表示最大值为c的方案个数，例如例子1中只有三个方案满足最大值为3，[3, 1], [3, 2] [3, 3]。

那么可以不可以定义一个函数直接就是问题的解（而不用累加呢？），可以定义$f(a,b,c)$表示最大值小于等于c的方案个数，那么

• $f(a,b,c)=f(a,b,c-1)+f(a-1,b-1,c-1)+(f(a-1,b,c)-f(a-1,b,c-1))\ast c$

接着思考边界条件，当$a=1\&b=1$，此时$f(1,1,c)=c$。

class Solution:
    def numOfArrays(self, n: int, m: int, k: int) -> int:
        mod = 10**9 + 7

        dp = [[[0] * (m + 1) for _ in range(k + 1)] for _ in range(n + 1)]
        for c in range(m + 1):
            dp[1][1][c] = c
        for a in range(2, n + 1):
            for b in range(1, k + 1):
                for c in range(b, m + 1):
                    dp[a][b][c] = (dp[a][b][c - 1] + dp[a - 1][b - 1][c - 1] + \
                                   (dp[a - 1][b][c] - dp[a - 1][b][c - 1] + mod) * c) % mod
        return dp[-1][-1][-1]

非常简洁优雅~

reference:

https://leetcode.com/problems/build-array-where-you-can-find-the-maximum-exactly-k-comparisons/discuss/591186/C%2B%2B-DP-solution-4ms-98-w-explanation

我将该问题的其他语言版本添加到了我的GitHub Leetcode

如有问题，希望大家指出！！！