Help with Atcoder TDPC (D-Dice)

Problem Link:

How I Solved the Problem: First of all I factored the number D as a product of primes 2, 3, 5 by storing the number of times each prime appears. If D has a prime factor other than 2,3,5 then the required probability will be 0. Otherwise there might exist some non-zero probability. For that I came up with the following recurrence.
prob(n,i,j,k) = (prob(n-1,i,j,k) + prob(n-1,i+1,j,k) + prob(n-1,i,j+1,k) + prob(n-1,i+2,j,k) + prob(n-1,i,j,k+1) + prob(n-1,i+1,j+1,k))/6. Here p(n,i,j,k) is the probability of getting a multiple of D when I have n throws remaining, with i 2’s, j 3’s, k 5’s accumulated till now.

My Code:

#include <bits/stdc++.h>
using namespace std;

#define FASTIO ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
using ll = long long;
using lld = long double;

const int MOD = 1e9+7;
const ll INF = 1e18;

int num;
ll d;

int cnt2 = 0, cnt3 = 0, cnt5 = 0;

void factor(int p, int& cnt) {
    while(d%p == 0) {
        d /= p;

lld dp[105][65][40][30];

lld prob(int n, int i, int j, int k) {
    if(n == 0 && i >= cnt2 && j >= cnt3 && k >= cnt5) return dp[n][i][j][k] = 1;
    if(n == 0 && (i <= cnt2 || j <= cnt3 || k <= cnt5)) return dp[n][i][j][k] = 0;
    if(dp[n][i][j][k] != -1) return dp[n][i][j][k];
    dp[n][i][j][k] = (prob(n-1,i,j,k) + prob(n-1,i+1,j,k) + prob(n-1,i,j+1,k)
                      + prob(n-1,i+2,j,k) + prob(n-1,i,j,k+1) + prob(n-1,i+1,j+1,k))/6;
    return dp[n][i][j][k];

int main() {
    cin >> num >> d;
    factor(2,cnt2); factor(3,cnt3); factor(5,cnt5);
    if(d != 1) {
        cout << 0 << "\n";
    else {
        for(int n=0; n<105; n++) {
            for(int i=0; i<65; i++) {
                for(int j=0; j<40; j++) {
                    for(int k=0; k<30; k++) {
                        dp[n][i][j][k] = -1;
        cout << fixed << setprecision(6) << prob(num,0,0,0) << "\n";
    return 0;

I have tried this code on the sample test cases and I am getting the correct answers but still the code gives a WA on submission. I tried looking for a similar problem (with the solution) but couldn’t find any. So I don’t know what is wrong with the logic or the implementation. Can anyone please help?


It is not showing in English…can you provide the link

You’ll have to use Translate. As far as I know it is not available in English.

I wrote iterative DP with the same states, but I used 3D dp to cut down on the memory. You can have a look to get a better idea. Here is my submission.

@nishitsharma03 Thanks for your response. The problem is that I am more used to memoization than iterative dp, though I am trying to get better at that. I saw your code but I am having a little difficulty in understanding how you managed to do the state transitions bottom-up with one less state in the dp table. Can u please expand upon that a little bit? And is it true that my code gets a WA due to memory and time issues, or was there something else wrong in my code?