LeetCode 1223. Dice Roll Simulation

原题链接在这里：https://leetcode.com/problems/dice-roll-simulation/

题目：

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

题解：

Let dp[i][j][k] denotes number of distinct sequences up to i dice, ending with number j with k repeating j at the end.

For previous p, if j != p, then dp[i][j][1] = sum(dp[i - 1][p][k]), sum of previous ending with p, with k repeating p at the end.

Else if j == p, we need to make sure k + 1 <= rollMax[j], dp[i][j][k + 1] = dp[i - 1][j][k]. otherwise it would be 0. 

Time Complexity: O(n*m*m*rMax). m = 6. rMax = max(rollMax[i]).

Space: O(n*m*rMax).

AC Java:

class Solution {
    public int dieSimulator(int n, int[] rollMax) {
        int mod = 1000000007;
        int rMax = 15;
        
        int [][][] dp = new int[n + 1][6][rMax + 1];
        for(int i = 0; i < 6; i++){
            dp[1][i][1] = 1;
        }
        
        for(int i = 2; i <= n; i++){
            for(int j = 0; j < 6; j++){
                for(int p = 0; p < 6; p++){
                    for(int k = 1; k <= rMax; k++){
                        if(j != p){
                            dp[i][j][1] = (dp[i][j][1] + dp[i - 1][p][k]) % mod;
                        }else if(k < rollMax[j]){ // same number, make sure k + 1 <= rollMax[j]
                            dp[i][j][k + 1] = dp[i - 1][j][k];
                        }
                    }
                }
            }
        }
        
        int res = 0;
        for(int j = 0; j < 6; j++){
            for(int k = 1; k <= rMax; k++){
                res = (res + dp[n][j][k]) % mod;
            }
        }
        
        return res;
    }
}