【多维DP】【hard】力扣1223. 掷骰子模拟

有一个骰子模拟器会每次投掷的时候生成一个 1 到 6 的随机数。

不过我们在使用它时有个约束，就是使得投掷骰子时，连续掷出数字 i 的次数不能超过 rollMax[i]（i 从 1 开始编号）。

现在，给你一个整数数组 rollMax 和一个整数 n，请你来计算掷 n 次骰子可得到的不同点数序列的数量。

假如两个序列中至少存在一个元素不同，就认为这两个序列是不同的。由于答案可能很大，所以请返回模 10^9 + 7 之后的结果。

示例 1：

输入：n = 2, rollMax = [1,1,2,2,2,3]

输出：34

解释：我们掷 2 次骰子，如果没有约束的话，共有 6 * 6 = 36 种可能的组合。但是根据 rollMax 数组，数字 1 和 2 最多连续出现一次，所以不会出现序列 (1,1) 和 (2,2)。因此，最终答案是 36-2 = 34。

示例 2：

输入：n = 2, rollMax = [1,1,1,1,1,1]

输出：30

示例 3：

输入：n = 3, rollMax = [1,1,1,2,2,3]

输出：181

提示：

1 <= n <= 5000

rollMax.length == 6

1 <= rollMax[i] <= 15

cpp 复制代码

class Solution {
public:
    static constexpr int mod = 1e9 + 7;
    int dieSimulator(int n, vector<int>& rollMax) {
        int m = *max_element(rollMax.begin(), rollMax.end());
        vector dp(n+1, vector(6, vector<int>(m+1)));
        for(int j = 0; j < 6; j++){
            dp[1][j][1] = 1;
        }

        for(int i = 2; i <= n; i++){
            //枚举已完成最后一次投掷的点数
            for(int j = 0; j < 6; j++){
                for(int k = 1; k <= rollMax[j]; k++){
                    //枚举这一次点数
                    for(int p = 0; p < 6; p++){
                        if(p != j){
                            dp[i][p][1] = (dp[i][p][1] + dp[i-1][j][k]) % mod;
                        }
                        else if(k + 1 <= rollMax[j]){
                            dp[i][p][k+1] = (dp[i][p][k+1] + dp[i-1][j][k]) % mod;
                        }
                    }
                }
            }
        }
        int res = 0;
            for(int j = 0; j < 6; j++){
                for(int k = 1; k <= rollMax[j]; k++){
                    res = (res + dp[n][j][k]) % mod;
                }
            }
            return res;
    }
};

时间复杂度：O(n m^2 k)，其中 n 是掷骰子的次数，m 是随机数的种类数，在本题中等于 6，k 是 rollMax 数组中的最大值。

空间复杂度：O(nmk)。

这道题我们定义一个三维数组dp[i][j][k]表示已经完成了i次投掷，并且第i次投掷的结果是是数字j，然后j已经连续摇到了k次。

那么我们就可以遍历i，然后开始遍历j，也就是枚举最后一次投掷的点数，然后再枚举k，代表j已经连续摇到了k次。那么接下来我们枚举p，p的含义是我们新投一次骰子摇到的数字是多少吗，当p不等于j的时候，那么列出状态转移方程 dp[i][p][1] = (dp[i][p][1] + dp[i-1][j][k]) % mod;，否则当p等于j的时候，并且k+1不大于j能连续的最高次数，列出状态转移方程dp[i][p][k+1] = (dp[i][p][k+1] + dp[i-1][j][k]) % mod;。

最后我们定义一个变量res，来记录以不同点数结尾并且连续进行了不同次数的dp的总和，最后返回res。