1223. 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

Rust Solution

struct Solution;

const M: i32 = 1_000_000_007;

impl Solution {
    fn die_simulator(n: i32, roll_max: Vec<i32>) -> i32 {
        let n = n as usize;
        let mut dp: Vec<Vec<Vec<i32>>> = vec![vec![vec![0; 5001]; 16]; 6];
        Self::dfs(n, -1, 1, &mut dp, &roll_max)
    }

    fn dfs(
        n: usize,
        prev: i32,
        repeat: usize,
        dp: &mut Vec<Vec<Vec<i32>>>,
        roll_max: &[i32],
    ) -> i32 {
        if n == 0 {
            1
        } else {
            if prev >= 0 && dp[prev as usize][repeat][n] > 0 {
                return dp[prev as usize][repeat][n];
            }
            let mut sum = 0;
            for i in 0..6 {
                if i == prev {
                    if repeat < roll_max[i as usize] as usize {
                        sum += Self::dfs(n - 1, i, repeat + 1, dp, roll_max);
                        sum %= M;
                    }
                } else {
                    sum += Self::dfs(n - 1, i, 1, dp, roll_max);
                    sum %= M;
                }
            }
            if prev >= 0 {
                dp[prev as usize][repeat][n] = sum;
            }
            sum
        }
    }
}

#[test]
fn test() {
    let n = 2;
    let roll_max = vec![1, 1, 2, 2, 2, 3];
    let res = 34;
    assert_eq!(Solution::die_simulator(n, roll_max), res);
    let n = 2;
    let roll_max = vec![1, 1, 1, 1, 1, 1];
    let res = 30;
    assert_eq!(Solution::die_simulator(n, roll_max), res);
    let n = 3;
    let roll_max = vec![1, 1, 1, 2, 2, 3];
    let res = 181;
    assert_eq!(Solution::die_simulator(n, roll_max), res);
    let n = 4;
    let roll_max = vec![2, 1, 1, 3, 3, 2];
    let res = 1082;
    assert_eq!(Solution::die_simulator(n, roll_max), res);
}

Having problems with this solution? Click here to submit an issue on github.