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

解法

方法一：记忆化搜索

思考

连续同一点数受 \(rollMax\) 限制，\(n\) 达 \(5000\)，直接枚举序列不可行。方案数只依赖「还剩几次、上次点数、该点已连续几次」。

记忆化 \(dfs(i,j,x)\)：第 \(i\) 次起、上一次为 \(j\) 且已连续 \(x\) 次。枚举下一点 \(k\)，不同则重置连续次数，相同则须 \(x\) 仍低于上限。初态 \((0,0,0)\) 表示尚未掷出。

状态 \(O(n\times 6\times 15)\)，转移枚举 \(6\) 个点数，记忆化避免重复子问题。

我们可以设计一个函数 \(dfs(i, j, x)\) 表示从第 \(i\) 次掷骰子开始，当前掷出的点数为 \(j\)，且连续掷出 \(j\) 的次数为 \(x\) 的方案数。其中 \(j\) 的取值范围为 \([1, 6]\)，而 \(x\) 的取值范围为 \([1, rollMax[j - 1]]\)。那么答案就是 \(dfs(0, 0, 0)\)。

函数 \(dfs(i, j, x)\) 的计算过程如下：

如果 \(i \ge n\)，说明已经掷完了 \(n\) 次骰子，返回 \(1\)。
否则，我们枚举下一次掷出的点数 \(k\)，如果 \(k \ne j\)，那么我们可以直接掷出 \(k\)，此时连续掷出 \(j\) 的次数 \(x\) 就会被重置为 \(1\)，因此方案数为 \(dfs(i + 1, k, 1)\)。如果 \(k = j\)，那么我们需要判断 \(x\) 是否小于 \(rollMax[j - 1]\)，如果小于，那么我们可以继续掷出 \(j\)，此时连续掷出 \(j\) 的次数 \(x\) 就会加 \(1\)，因此方案数为 \(dfs(i + 1, j, x + 1)\)。最后将所有方案数相加，即为 \(dfs(i, j, x)\) 的值。注意答案可能很大，因此需要对 \(10^9 + 7\) 取模。

过程中，我们可以使用记忆化搜索避免重复计算。

时间复杂度 \(O(n \times k^2 \times M)\)，空间复杂度 \(O(n \times k \times M)\)。其中 \(k\) 为点数的取值范围，而 \(M\) 为连续掷出某个点数的最大次数。

Python3JavaC++Go

class Solution:
    def dieSimulator(self, n: int, rollMax: List[int]) -> int:
        @cache
        def dfs(i, j, x):
            if i >= n:
                return 1
            ans = 0
            for k in range(1, 7):
                if k != j:
                    ans += dfs(i + 1, k, 1)
                elif x < rollMax[j - 1]:
                    ans += dfs(i + 1, j, x + 1)
            return ans % (10**9 + 7)

        return dfs(0, 0, 0)

class Solution {
    private Integer[][][] f;
    private int[] rollMax;

    public int dieSimulator(int n, int[] rollMax) {
        f = new Integer[n][7][16];
        this.rollMax = rollMax;
        return dfs(0, 0, 0);
    }

    private int dfs(int i, int j, int x) {
        if (i >= f.length) {
            return 1;
        }
        if (f[i][j][x] != null) {
            return f[i][j][x];
        }
        long ans = 0;
        for (int k = 1; k <= 6; ++k) {
            if (k != j) {
                ans += dfs(i + 1, k, 1);
            } else if (x < rollMax[j - 1]) {
                ans += dfs(i + 1, j, x + 1);
            }
        }
        ans %= 1000000007;
        return f[i][j][x] = (int) ans;
    }
}

class Solution {
public:
    int dieSimulator(int n, vector<int>& rollMax) {
        int f[n][7][16];
        memset(f, 0, sizeof f);
        const int mod = 1e9 + 7;
        function<int(int, int, int)> dfs = [&](int i, int j, int x) -> int {
            if (i >= n) {
                return 1;
            }
            if (f[i][j][x]) {
                return f[i][j][x];
            }
            long ans = 0;
            for (int k = 1; k <= 6; ++k) {
                if (k != j) {
                    ans += dfs(i + 1, k, 1);
                } else if (x < rollMax[j - 1]) {
                    ans += dfs(i + 1, j, x + 1);
                }
            }
            ans %= mod;
            return f[i][j][x] = ans;
        };
        return dfs(0, 0, 0);
    }
};

func dieSimulator(n int, rollMax []int) int {
    f := make([][7][16]int, n)
    const mod = 1e9 + 7
    var dfs func(i, j, x int) int
    dfs = func(i, j, x int) int {
        if i >= n {
            return 1
        }
        if f[i][j][x] != 0 {
            return f[i][j][x]
        }
        ans := 0
        for k := 1; k <= 6; k++ {
            if k != j {
                ans += dfs(i+1, k, 1)
            } else if x < rollMax[j-1] {
                ans += dfs(i+1, j, x+1)
            }
        }
        f[i][j][x] = ans % mod
        return f[i][j][x]
    }
    return dfs(0, 0, 0)
}

方法二：动态规划

思考

记忆化是按递归展开的同一转移。改成自底向上：\(f[i][j][x]\) 表示掷完 \(i\) 次、结尾为 \(j\) 且连续 \(x\) 次的方案，由 \(i-1\) 层按相同规则累加，去掉递归栈，便于改成滚动数组。

我们可以将方法一中的记忆化搜索改为动态规划。

定义 \(f[i][j][x]\) 表示投掷前 \(i\) 次骰子，且第 \(i\) 次投掷的点数为 \(j\)，且连续投掷点数 \(j\) 的次数为 \(x\) 的方案数。初始时 \(f[1][j][1] = 1\)，其中 \(1 \leq j \leq 6\)。答案即是：

\[ \sum_{j=1}^6 \sum_{x=1}^{rollMax[j-1]} f[n][j][x] \]

我们枚举上一次投掷的点数为 \(j\)，且连续投掷点数 \(j\) 的次数为 \(x\)，那么当前投掷的点数可以为 \(1, 2, \cdots, 6\)，如果当前投掷的点数为 \(k\)，那么有如下两种情况：

如果 \(k \neq j\)，那么我们可以直接投掷出 \(k\)，此时连续投掷点数 \(j\) 的次数 \(x\) 就会被重置为 \(1\)，因此方案数 \(f[i][k][1]\) 就会增加 \(f[i-1][j][x]\)。
如果 \(k = j\)，那么我们需要判断 \(x+1\) 是否小于等于 \(rollMax[j-1]\)，如果小于等于，那么我们可以继续投掷出 \(j\)，此时连续投掷点数 \(j\) 的次数 \(x\) 就会加 \(1\)，因此方案数 \(f[i][j][x+1]\) 就会增加 \(f[i-1][j][x]\)。

最终的答案即为所有 \(f[n][j][x]\) 的和。