title

problem link

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

solution

DP[i][j] means the total number of combinations at round i we get a j at current round, considering the constraint.

Explanation Let J be any integer in 1 to 6, and the maximum roll amount be k, then at round i we only allow: J^CJ, J^CJJ, ..., J^CJJ...J, (k consecutive J) where J^C represents a number that is not J, namely the complementary set of J.

DP approach It is not difficult to see that DP[i][j] = DP[i-k][j^C] + ... + DP[i-1][j^C]. Upon noticing that DP[i][j^c] = sum[i] - DP[i][j], where sum[i] represents total number of combinations at round i considering the constraint, we end up with a two-dimensional DP, and an additional one-dimensional DP to quickly obtain sum[i].

The code should be pretty self-explanatory.

class Solution:
    def dieSimulator(self, n: int, rollMax: List[int]) -> int:
        p = 10**9+7
        dp = [[0 for _ in range(6)] for _ in range(n + 1)]
        sum = [0 for _ in range(n + 1)]
        sum[0] = 1
        # dp i, j . res for roll i times when last with face j
        for i in range(1, n + 1):
            for j in range(6):
                for k in range(1, rollMax[j] + 1):
                    if i >= k:
                        dp[i][j] += sum[i - k] - dp[i - k][j]
                        dp[i][j] %= p
                sum[i] += dp[i][j]
        return sum[n] % p