DIGIT DP Tutorial [ITERATIVE]

№	Пользователь	Рейтинг
1	tourist	3993
2	jiangly	3743
3	orzdevinwang	3707
4	Radewoosh	3627
5	jqdai0815	3620
6	Benq	3564
7	Kevin114514	3443
8	ksun48	3434
9	Rewinding	3397
10	Um_nik	3396

№	Пользователь	Вклад
1	cry	167
2	Um_nik	163
3	maomao90	162
3	atcoder_official	162
5	adamant	159
6	-is-this-fft-	158
7	awoo	155
8	TheScrasse	154
9	Dominater069	153
10	djm03178	152

Why iterative?

I learned digit DP few years back but only recently I realised that the recursive solution is sometimes hard to debug and difficult to reason about. So this article aims to provide an iterative solution template that can be extended to solve similar problems.

Problem Statement

Find count of numbers in range [L, R] such that sum of its digits is a prime number.

$$$1 <= L,R <= 10^{18}$$$

Basic Idea

A brute force approach would be to iterate through each number from L to R and check if sum of its digits is prime. This will obviously timeout as constraints are too large.
So which particular property can we count such that we don't consume too much time/memory?
Sum of digits of numbers! Sum of digits of a number cannot exceed 180 for 18 digits, thus we will have to track only 180 states which will save time and memory.This is the key idea behind most digit DP problems: identify and track a property which is finite and will help us reach the answer.
We will try to create a function $$$f(x)$$$ which returns all good numbers from $$$[0,x]$$$. Later we can use $$$f(R) - f(L-1)$$$ to find our answer.

Detailed Explanation

Let's declare our dp as follows:

    dp[20][2][200]

20 $$$\rightarrow$$$ maximum number of digits that our dp will support (18 to be precise)
2 $$$\rightarrow$$$ tight condition (explained later)
200 $$$\rightarrow$$$ maximum possible sum of digits of a number

For better understanding, try to think of indexes as $$$dp[i][tight][sum]$$$

What does $$$dp[i][tight][sum]$$$ even mean?

$$$dp[i][0][sum] \rightarrow$$$ count of suffixes that can be formed starting from index i, whose digits add up to $$$sum$$$

$$$dp[i][1][sum] \rightarrow$$$ count of suffixes that can be formed starting from index i, whose digits add up to $$$sum$$$
such that the formed suffix is not greater than corresponding suffix in input string

Base Cases dp[n][0][0] = dp[n][1][0] = 1
There exists $$$1$$$ empty suffix with sum of its digits $$$=0$$$
We will move from right to left while prepending digits at every index till we finally calculate $$$dp[0][..][..]$$$ which denotes the entire input string

What is $$$tight$$$?

For every state of dp we also need to know if the current suffix formed includes all unbounded numbers or only the numbers less than or equal to $$$suffix_{input}$$$. This information is required because when we are prepending digits to build $$$dp[i]$$$ from $$$dp[i+1]$$$, we will have to choose between bounded/unbounded suffixes based on our current digit. Try to understand this based on the following recurrence relations.
We can break $$$dp[i][tight][sum]$$$ into subproblems as follows:

$$$dp[i][0][sum] = \sum\limits_{d=0}^{9} dp[i+1][0][sum-d] $$$

$$$dp[i][1][sum] = dp[i+1][1][sum-ss[i]] + \sum\limits_{d=0}^{ss[i]-1} dp[i+1][0][sum-d] $$$

Code

int digit_dp(string ss) {
    int n = ss.size();
 
    //empty suffixes having sum=0
    dp[n][0][0] = 1;
    dp[n][1][0] = 1;
 
    for(int i = n-1; i >=0 ; i--) {
        for(int tight = 0; tight < 2 ; tight++) {
            for(int sum = 0; sum < 200 ; sum++) {
                if(tight) {
                    for(int d = 0; d <= ss[i] - '0' ; d++) {
                        dp[i][1][sum] += (d == ss[i]-'0') ? dp[i+1][1][sum-d] : dp[i+1][0][sum-d];
                    }
                }
                else {
                    for(int d = 0; d < 10 ; d++) {
                        dp[i][0][sum] += dp[i+1][0][sum-d];
                    }
                }
            }
        }
    }
    int ans = 0;
    for(int i = 0; i < 200; i++) {
        if(isPrime(i))
        ans += dp[0][1][i];
    }
    return ans;
}

Complexity

As obvious from the loops: $$$O(digits*2*sum*10)$$$

End Notes

Find above SPOJ problem here: GONE

Similar SPOJ problem where you can extend above solution: RAONE
Recent codeforces problem which requires a similar approach: https://codeforces.me/contest/1341/problem/D
Educational Codeforces Digit DP problem: https://codeforces.me/contest/1036/problem/C
This is my first article on codeforces. All comments/criticism are welcome.

Solution GONE

#include<bits/stdc++.h>
using namespace std;
 
int dp[20][2][200];
int primes[46] = {2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97, 101, 103, 107, 109, 113, 127, 131, 137, 139, 149, 151, 157, 163, 167, 173, 179, 181, 191, 193, 197, 199};
/////////////////////////////////////////////////
 
bool isPrime(int x) {
    for (int i = 0; i < 46; i++)
        if (primes[i] == x)
            return true;
    return false;
}
 
//count of all numbers in [0,ss] which have prime sum of digits
//dp[i][tight][sum]  =>  number of suffixes[i..n-1] with sum of its digits = 'sum'
//NOTE: prime is only checked at last, dp itself only holds count of numbers
 
int digit_dp(string ss) {
    int n = ss.size();
 
    //empty suffixes having sum=0
    dp[n][0][0] = 1;
    dp[n][1][0] = 1;
 
    for(int i = n-1; i >=0 ; i--) {
        for(int tight = 0; tight < 2 ; tight++) {
            for(int sum = 0; sum < 200 ; sum++) {
                if(tight) {
                    for(int d = 0; d <= ss[i] - '0' ; d++) {
                        dp[i][1][sum] += (d == ss[i]-'0') ? dp[i+1][1][sum-d] : dp[i+1][0][sum-d];
                    }
                }
                else {
                    for(int d = 0; d < 10 ; d++) {
                        dp[i][0][sum] += dp[i+1][0][sum-d];
                    }
                }
            }
        }
    }
    int ans = 0;
    for(int i = 0; i < 200; i++) {
        if(isPrime(i))
        ans += dp[0][1][i];
    }
    return ans;
}
 
void reset_dp() {
    for(int i = 0; i < 20; i++) 
        for(int j = 0; j < 2; j++) 
            for(int k = 0; k < 200; k++) 
                dp[i][j][k] = 0;
}
 
void purple()
{
    int a, b;
    cin >> a >> b;
    int countA = digit_dp(to_string(a - 1));
    reset_dp();
    int countB = digit_dp(to_string(b));
    reset_dp();
    cout << countB - countA << endl;
}
 
signed main() {
    int t=1;
    cin>>t;
    while(t--) purple();
}

Solution RAONE

#include<bits/stdc++.h>
using namespace std;
 
int dp[12][2][50][50]; // {idx, tight, odd, even}
 
int digit_dp(string ss) {
    int n = ss.size();
 
    //empty suffixes having sum=0
    dp[n][0][0][0] = 1;
    dp[n][1][0][0] = 1;
 
    for(int i = n-1; i >=0 ; i--) {
        for(int tight = 0; tight < 2 ; tight++) {
            for(int odd = 0; odd < 50 ; odd++) {
                for(int even = 0; even < 50 ; even++) {
                    if((n - i) % 2) {
                        if(tight) {
                            for(int d = 0; d <= ss[i]-'0'; d++) {
                                dp[i][1][odd][even] += (d == ss[i]-'0') ? dp[i+1][1][odd-d][even] : dp[i+1][0][odd-d][even] ;
                            }
                        }
                        else {
                            for(int d = 0; d < 10; d++) {
                                dp[i][0][odd][even] += dp[i+1][0][odd-d][even] ;
                            }
                        }
                    }
                    else {
                        if(tight) {
                            for(int d = 0; d <= ss[i]-'0'; d++) {
                                dp[i][1][odd][even] += (d == ss[i]-'0') ? dp[i+1][1][odd][even-d] : dp[i+1][0][odd][even-d] ;
                            }
                        }
                        else {
                            for(int d = 0; d < 10; d++) {
                                dp[i][0][odd][even] +=  dp[i+1][0][odd][even-d] ;
                            }
                        }
                    }   
                }
            }
        }
    }
    int ans = 0;
    for(int o = 0; o < 50; o++) {
        for(int e = 0; e < 50; e++) {
            if(e-o==1)
            ans += dp[0][1][o][e];
        } 
    }
    return ans;
}
 
void reset_dp() {
    for(int i = 0; i < 12; i++) {
        for(int j = 0; j < 2; j++) {
            for(int k = 0; k < 50; k++) {
                for(int l = 0; l < 50; l++) {
                    dp[i][j][k][l] = 0;
                }
            }
        }
    }
}
 
void purple()
{
    int a, b;
    cin >> a >> b;
    int countA = digit_dp(to_string(a - 1));
    reset_dp();
    int countB = digit_dp(to_string(b));
    reset_dp();
    cout << countB - countA << endl;
}
 
signed main() {
    int t=1;
    cin>>t;
    while(t--) purple();
}

Solution Problem D

#include<bits/stdc++.h>
using namespace std;
// #define int long long int
#define vec vector
#define pb push_back
#define s(x) scanf("%d",&x)
#define all(x)	x.begin(),x.end()
#define pii pair<int,int>
#define fi first
#define se second
#define mod 1000000007
#define SZ 200005
#define endl "\n"
#define trace(x) cerr << #x << ": " << x << endl
#define boost ios::sync_with_stdio(0);cin.tie(0)

int n,k;
 
vector<string> dig = {"1110111", "0010010", "1011101", "1011011", "0111010", "1101011", "1101111", "1010010", "1111111", "1111011"};

void purple()
{
    cin >> n >> k;
    string a[n];
    for(int i = 0; i < n; i++) cin >> a[i];

    vector< vector<int> > dist(n, vector<int>(10));

    for(int i = 0; i < n; i++) {
        for(int d = 0; d < 10; d++) {
            for(int j = 0; j < 7; j++) {
                if(a[i][j]=='1' and dig[d][j]=='0'){
                    dist[i][d] = -1;
                    break;
                }
                else if(a[i][j]!=dig[d][j])
                dist[i][d]++;
            }
        }
    }

    vector< vector<bool> > dp(n+1, vector<bool> (k+1, false));
    dp[n][0] = true;

    for(int i = n-1; i >=0 ; i--) {
        for(int j = 0; j <=k ; j++) {
            for(int d = 0; d < 10; d++) {
                dp[i][j] = dp[i][j] || (dist[i][d]!=-1 && j-dist[i][d]>=0 && dp[i+1][j-dist[i][d]]);
            }
        }
    }
    if(!dp[0][k]) {cout << "-1\n"; return;}

    for(int i = 0; i < n; i++) {
        for(int d = 9; d >= 0; d--) {
            if(dist[i][d]!=-1 && k-dist[i][d]>=0 && dp[i+1][k-dist[i][d]]) {
                cout << d;
                k -= dist[i][d];
                break;
            }
        }
    }
}

signed main() {
    boost;
    int t=1;
    // cin>>t;
    while(t--) purple();
}