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.
Two sequences are considered different if at least one element differs from each other. Since the answer may be too large, return it modulo 10^9 + 7
.
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 <= 5000
rollMax.length == 6
1 <= rollMax[i] <= 15
思路:dp[i][j][k]表示使用了i次,最后以j结尾,并且j重复了k次
from collections import defaultdict
class Solution(object):
def dieSimulator(self, n, rollMax):
"""
:type n: int
:type rollMax: List[int]
:rtype: int
"""
mod = 10**9+7
dp = [[[0 for _ in range(16)] for _ in range(6)] for _ in range(n)]
for j in range(6): dp[0][j][1]=1
for i in range(1,n):
for j in range(6):
for k in range(1,rollMax[j]+1):
for jj in range(6):
if j==jj:
if k+1<=rollMax[j]:
dp[i][jj][k+1] += dp[i-1][j][k]
else:
dp[i][jj][1] += dp[i-1][j][k]
return sum(sum(row) for row in dp[-1])%mod