1223. 掷骰子模拟

ITCharge大约 5 分钟

---

1223. 掷骰子模拟

标签：数组、动态规划
难度：困难

题目大意

描述：有一个骰子模拟器，每次投掷会生成一个 $1$ 到 $6$ 的随机数。但有个约束：连续掷出数字 $i$ 的次数不能超过 $rollMax[i]$ （ $i$ 从 $1$ 开始编号）。给定一个整数数组 $rollMax$ 和一个整数 $n$ 。

要求：计算掷 $n$ 次骰子可得到的不同点数序列的数量。答案可能很大，返回模 $10^9 + 7$ 之后的结果。

说明：

$1 \le n \le 5000$ 。
$rollMax.length == 6$ 。
$1 \le rollMax[i] \le 15$ 。

示例：

示例 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

解题思路

思路 1：动态规划

1. 阶段划分

按照投掷次数 $i$ 划分阶段。每次投掷时，我们需要知道上一次掷出的数字以及它连续出现了几次，才能判断当前能掷什么数字。所以状态需要包含三个维度：投掷次数、当前数字、当前数字的连续出现次数。

2. 定义状态

定义 $dp[i][j][k]$ 表示掷了 $i$ 次骰子，且第 $i$ 次掷出的数字是 $j$ （ $j$ 从 $0$ 到 $5$ ，对应数字 $1$ 到 $6$ ），并且 $j$ 在末尾连续出现了 $k$ 次的不同序列数量。

3. 状态转移方程

初始状态（ $i = 1$ ）：第一次掷骰子，数字 $j$ 出现 $1$ 次，每种情况有 $1$ 种序列。

dp[1][j][1] = 1

当 $i > 1$ 时，分两种情况：

情况 A：当前数字与前一次不同（ $k = 1$ ）

当前掷出 $j$ 且只连续 $1$ 次，说明上一次掷出的不是 $j$ 。上一次可以是其他 $5$ 个数字中的任意一个，且连续次数不限（只要不超过各自的 $rollMax$ 限制）。所以：

dp[i][j][1] = \sum_{j' \neq j} \sum_{k'=1}^{rollMax[j']} dp[i-1][j'][k']

可以理解为：掷 $i-1$ 次的总序列数减去第 $i-1$ 次以 $j$ 结尾的所有序列数。

dp[i][j][1] = total[i-1] - \sum_{k'=1}^{rollMax[j]} dp[i-1][j][k']

其中 $total[i-1]$ 表示掷 $i-1$ 次的所有合法序列数。

情况 B：当前数字与前一次相同（ $k > 1$ ）

当前掷出 $j$ 且连续 $k$ 次，说明上一次也掷出了 $j$ 且连续了 $k-1$ 次（前提是 $k \le rollMax[j]$ ）：

dp[i][j][k] = dp[i-1][j][k-1]

4. 初始条件

$dp[1][j][1] = 1$ ，其中 $j = 0,1,...,5$ 。
$total[1] = 6$ 。

5. 最终结果

ans = \sum_{j=0}^{5} \sum_{k=1}^{rollMax[j]} dp[n][j][k] \mod (10^9 + 7)

也就是 $total[n]$ 。

结合示例 1 走一遍：

$n = 2, rollMax = [1,1,2,2,2,3]$

初始化：

$dp[1][0][1] = dp[1][1][1] = ... = dp[1][5][1] = 1$
$total[1] = 6$

$i = 2$ ：

$j=0$ （数字 $1$ ）： $\sum_{k} dp[1][0][k] = 1$ ， $dp[2][0][1] = total[1] - 1 = 5$ ； $rollMax[0]=1$ ，没有 $k>1$ 的情况。
$j=1$ （数字 $2$ ）：同理 $dp[2][1][1] = 5$ 。
$j=2$ （数字 $3$ ）： $\sum_{k} dp[1][2][k] = 1$ ， $dp[2][2][1] = 5$ ； $rollMax[2]=2$ ， $dp[2][2][2] = dp[1][2][1] = 1$ 。
$j=3,4,5$ ：与 $j=2$ 类似， $rollMax$ 分别为 $2,2,3$ ，所以 $dp[2][j][1]=5$ ， $dp[2][j][2]=1$ （对 $j=5$ 还有 $dp[2][5][2]=1$ ）。
$total[2] = 各 dp 之和 = 5 \times 6 + 1 \times 5 = 35$ ？

等等，应该直接按公式算更准确。让我们按递推：

$total[2] = \sum_j \sum_k dp[2][j][k]$

对于 $j=0$ （ $rollMax=1$ ）：只有 $k=1$ ， $dp=5$
对于 $j=1$ （ $rollMax=1$ ）：只有 $k=1$ ， $dp=5$
对于 $j=2$ （ $rollMax=2$ ）： $k=1: 5$ ， $k=2: 1$
对于 $j=3$ （ $rollMax=2$ ）： $k=1: 5$ ， $k=2: 1$
对于 $j=4$ （ $rollMax=2$ ）： $k=1: 5$ ， $k=2: 1$
对于 $j=5$ （ $rollMax=3$ ）： $k=1: 5$ ， $k=2: 1$

$total[2] = 5+5+5+1+5+1+5+1+5+1 = 34$ ，与示例一致。

思路 1：代码

class Solution:
    def dieSimulator(self, n: int, rollMax: List[int]) -> int:
        MOD = 10 ** 9 + 7
        # dp[i][j][k]：掷 i 次，最后一次是数字 j，且 j 连续出现 k 次
        dp = [[[0] * (rollMax[j] + 1) for j in range(6)] for _ in range(n + 1)]
        # total[i]：掷 i 次的总序列数
        total = [0] * (n + 1)

        # 初始化：掷 1 次时，每个数字出现 1 次，共 6 种序列
        for j in range(6):
            dp[1][j][1] = 1
        total[1] = 6

        # 从第 2 次开始递推
        for i in range(2, n + 1):
            for j in range(6):
                # 情况 A：k = 1，本次数字与上次不同
                # 用总序列数减去上次以 j 结尾的序列数
                sum_j = sum(dp[i-1][j][k] for k in range(1, rollMax[j] + 1))
                dp[i][j][1] = (total[i-1] - sum_j) % MOD

                # 情况 B：k > 1，本次数字与上次相同，连续次数 +1
                for k in range(2, rollMax[j] + 1):
                    dp[i][j][k] = dp[i-1][j][k-1]

            # 计算本轮的总序列数
            total[i] = sum(dp[i][j][k] for j in range(6) for k in range(1, rollMax[j] + 1)) % MOD

        return total[n]

思路 1：复杂度分析

时间复杂度： $O(n \times 6 \times \max(rollMax))$ 。 $n \le 5000$ ，而 $rollMax$ 最大为 $15$ ，所以计算量约为 $5000 \times 6 \times 15 = 450000$ 次操作，完全可行。
空间复杂度： $O(n \times 6 \times \max(rollMax))$ 。可以优化为 $O(6 \times \max(rollMax))$ 使用滚动数组（因为 $dp[i]$ 只依赖 $dp[i-1]$ ），但 $n=5000$ 时原始三维数组的内存也完全可以接受。