Dice Roll Simulation LeetCode solution in Python

1
class Solution:
2
    def dieSimulator(self, n: int, rollMax: List[int]) -> int:
3
        dp = {}
4

5
        def solve(n, last, count):
6
            if n == 0:
7
                return 1
8
            if (n, last, count) in dp:
9
                return dp[(n, last, count)]
10
            ans = 0
11
            for i in range(6):
12
                if last == i:
13
                    if count == rollMax[i]:
14
                        continue
15
                    ans += solve(n - 1, last, count + 1)
16
                else:
17
                    ans += solve(n - 1, i, 1)
18
            dp[(n, last, count)] = ans
19
            return ans
20

21
        return solve(n, None, 0) % 1000000007