题目描述

有一个骰子模拟器会每次投掷的时候生成一个 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。

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

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

限制

• 1 <= n <= 5000
• rollMax.length == 6
• 1 <= rollMax[i] <= 15

算法

(动态规划) $O(nm)$

1. $f(i, j, k)$ 表示前 $i$ 次投，最后一次结果为 $j$ 且最后的数字连续了 $k$ 次的方案数。
2. 初始时，$f(0, j, 1) = 1$，其余均为 0。
3. 转移时，对于每一个 $i$ 和 $j$，枚举上一次最后的结果 $t$，如果 $j == t$，则转移 $f(i, j, k) = f(i, j, k) + f(i - 1, j, k - 1)$；否则 $f(i, j, 1) = f(i, j, 1) + f(i - 1, t, k)$。
4. 最终答案为 $f(n - 1, j, k)$ 的总和。

时间复杂度

• 状态数为 $O(nm)$，转移仅需要常数，故空间复杂度为 $O(nm)$，这里的 $m$ 为最大的 rollMax。

空间复杂度

• 状态数为 $O(nm)$， 故空间复杂度为 $O(nm)$。

C++ 代码

class Solution {
public:
    int dieSimulator(int n, vector<int>& rollMax) {
        const int mod = 1000000007;
        vector<vector<vector<int>>> f(n, vector<vector<int>>(6, vector<int>(16, 0)));
        for (int j = 0; j < 6; j++)
            f[0][j][1] = 1;

        for (int i = 1; i < n; i++)
            for (int j = 0; j < 6; j++)
                for (int t = 0; t < 6; t++) {
                    if (j == t) {
                        for (int k = 2; k <= rollMax[j]; k++)
                            f[i][j][k] = (f[i][j][k] + f[i - 1][j][k - 1]) % mod;
                    } else {
                        for (int k = 1; k <= rollMax[t]; k++)
                            f[i][j][1] = (f[i][j][1] + f[i - 1][t][k]) % mod;
                    }
                }

        int ans = 0;
        for (int j = 0; j < 6; j++)
            for (int k = 1; k <= rollMax[j]; k++)
                ans = (ans + f[n - 1][j][k]) % mod;

        return ans;
    }
};