leecode.1223. 掷骰子模拟

题目

掷骰子模拟

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

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

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

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

示例

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

思路分析

• 令dp[i][j][k]表示前i次，以j结尾，j重复k次的结果。那么对于dp[i][j][1]来说，因为这个数字新出现，所以等于前面一个数字以非j结尾。
• 那么对于dp[i][j]k来说，因为这个数字不是新出现，所以等于前面一个数字以j结尾，其数值进行累加dp[i - 1][j][k - 1]。

代码

class Solution {
public:
    int dieSimulator(int n, vector<int>& rollMax) {
        const int mod = 1e9 + 7;
    // dp[i][j][k]表示前i次，以j结尾，j重复k次的结果
        int dp[5005][6][18] = {0};
        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 cur = 0;cur < 6;cur++){//初始化dp[i][j][1]
                    if(cur == j) continue;
                    for(int k = 1;k <= rollMax[cur];k++)
                        dp[i][j][1] = (dp[i][j][1] + dp[i - 1][cur][k]) % mod;
                }
                for(int k = 2;k <= rollMax[j];k++){//计算dp[i][j][k](k > 1)
                    dp[i][j][k] = (dp[i][j][k] + dp[i - 1][j][k - 1]) % mod;   
                }    
            }
        }
        int ans = 0;
        for(int j = 0;j < 6;j++){//最后的结果就是累加和
            for(int k = 1;k <= rollMax[j];k++){
                ans = (ans + dp[n][j][k]) % mod; 
            }
        }
        return ans;
    }
};