Description
A die simulator generates a random number from 1 to 6 for each roll. You introduced a constraint to the generator such that it cannot roll the number i more than rollMax[i] (1-indexed) consecutive times.
Given an array of integers rollMax and an integer n, return the number of distinct sequences that can be obtained with exact n rolls. Since the answer may be too large, return it modulo 109 + 7.
Two sequences are considered different if at least one element differs from each other.
Example 1:
Input: n = 2, rollMax = [1,1,2,2,2,3] Output: 34 Explanation: There will be 2 rolls of die, if there are no constraints on the die, there are 6 * 6 = 36 possible combinations. In this case, looking at rollMax array, the numbers 1 and 2 appear at most once consecutively, therefore sequences (1,1) and (2,2) cannot occur, so the final answer is 36-2 = 34.
Example 2:
Input: n = 2, rollMax = [1,1,1,1,1,1] Output: 30
Example 3:
Input: n = 3, rollMax = [1,1,1,2,2,3] Output: 181
Constraints:
1 <= n <= 5000rollMax.length == 61 <= rollMax[i] <= 15
Solutions
Solution 1: Memoization Search
We can design a function dfs(i, j, x) to represent the number of schemes starting from the i-th dice roll, with the current dice roll being j, and the number of consecutive times j is rolled being x. The range of j is [1, 6], and the range of x is [1, rollMax[j - 1]]. The answer is dfs(0, 0, 0).
The calculation process of the function dfs(i, j, x) is as follows:
- If i \ge n, it means that n dice have been rolled, return 1.
- Otherwise, we enumerate the number k rolled next time. If k \ne j, we can directly roll k, and the number of consecutive times j is rolled will be reset to 1, so the number of schemes is dfs(i + 1, k, 1). If k = j, we need to judge whether x is less than rollMax[j - 1]. If it is less, we can continue to roll j, and the number of consecutive times j is rolled will increase by 1, so the number of schemes is dfs(i + 1, j, x + 1). Finally, add all the scheme numbers to get the value of dfs(i, j, x). Note that the answer may be very large, so we need to take the modulus of 109 + 7.
During the process, we can use memoization search to avoid repeated calculations.
The time complexity is O(n × k2 × M), and the space complexity is O(n × k × M). Here, k is the range of dice points, and M is the maximum number of times a certain point can be rolled consecutively.
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)(code-box)
Solution 2: Dynamic Programming
We can change the memoization search in Solution 1 to dynamic programming.
Define f[i][j][x] as the number of schemes for the first i dice rolls, with the i-th dice roll being j, and the number of consecutive times j is rolled being x. Initially, f[1][j][1] = 1, where 1 ≤ j ≤ 6. The answer is:
We enumerate the last dice roll as j, and the number of consecutive times j is rolled as x. The current dice roll can be 1, 2, …, 6. If the current dice roll is k, there are two cases:
- If k ≠ j, we can directly roll k, and the number of consecutive times j is rolled will be reset to 1. Therefore, the number of schemes f[i][k][1] will increase by f[i-1][j][x].
- If k = j, we need to judge whether x+1 is less than or equal to rollMax[j-1]. If it is less than or equal to, we can continue to roll j, and the number of consecutive times j is rolled will increase by 1. Therefore, the number of schemes f[i][j][x+1] will increase by f[i-1][j][x].
The final answer is the sum of all f[n][j][x].
The time complexity is O(n × k2 × M), and the space complexity is O(n × k × M). Here, k is the range of dice points, and M is the maximum number of times a certain point can be rolled consecutively.
class Solution: def dieSimulator(self, n: int, rollMax: List[int]) -> int: f = [[[0] * 16 for _ in range(7)] for _ in range(n + 1)] for j in range(1, 7): f[1][j][1] = 1 for i in range(2, n + 1): for j in range(1, 7): for x in range(1, rollMax[j - 1] + 1): for k in range(1, 7): if k != j: f[i][k][1] += f[i - 1][j][x] elif x + 1 <= rollMax[j - 1]: f[i][j][x + 1] += f[i - 1][j][x] mod = 10**9 + 7 ans = 0 for j in range(1, 7): for x in range(1, rollMax[j - 1] + 1): ans = (ans + f[n][j][x]) % mod return ans(code-box)
