Hello Codeforces,
Digit dp was one of the techniques which I had the hardest time trying to learn. I have seen multiple blogs on iterative digit dp. In this blog, I would like to share my recursive implementation of digit dp.

Problem Structure

Typically, digit dp is used in problems of finding the number of values in a range $$$[l,r]$$$ which follow a certain property. Let us call this $$$f(l,r)$$$. The following method will work only when $$$f(l,r) = f(0,r) - f(0,l-1)$$$, which is true for most of the problems. We will now focus on how to calculate $$$g(n) = f(0,n)$$$.

Idea behind the algorithm

To count all numbers less than or equal to $$$n$$$ which follow a certain property, we first count all numbers with the same number of digits as $$$n$$$. Here the numerical order will be the same as lexicographical order. We can now iterate over the common prefix of the two numbers and the first character which differs after the prefix. This gives us the following template.

int g(string s){
	int n = s.size();
	
	int ans = 0;
	for(int i = 0; i < n; i++){ // i is the first position of difference

		for(int j = (i == 0); j < s[i] - '0'; j++){ // keeping digit j at position i

			ans += valid_strings(n-i-1); // number of valid strings of length n-i-1

		}
	}
	
	if(valid(s)){
		ans += 1;
	}
	
	if(n > 1){ // All numbers with less digits
		ans += g(string(n-1,'9'));
	}
	
	return ans;
}

Some remarks are as follows.

• Adding memoisation to this recursive function gives polylog complexity in most cases.
• $$$j$$$ is initialised to $$$1$$$ for $$$i == 0$$$ since our numbers can't have leading zeroes.
• $$$\text{valid_strings}(\cdots)$$$ may be a dp table instead of a mathematical function. It might also depend on the prefix as well as the leftover length.

Some examples

Consider CSES Counting Numbers. Here we have to find the number of numbers less than or equal to $$$n$$$ that no two consecutive digits are the same. Here I get the following function. Full solution.

map<string,int> mp;
int g(string s){
	if(mp.find(s) != mp.end()){
		return mp[s];
	}
	
	int n = s.size();
	int &ans = mp[s];
	
	bool ok = true;
	
	for(int i = 0; i < n; i++){
		
		for(int j = (i == 0); j < s[i] - '0'; j++){
			if(i == 0 || s[i-1] != j + '0'){
			       ans += power(9,n-i-1);
                        }
		}
		
		if(i != 0 && s[i] == s[i-1]){
			ok = false;
			break;
		}
	}
	
	ans += ok;
	
	if(n > 1){
		ans += get(string(n-1,'9'));
	}
	
	return ans;
}

CF 919B is an example where a dp table is needed to implement the $$$\text{valid_strings}(\cdots)$$$ function. There is a preparatory phase for filling the dp table. The relevant functions are provided below. Full solution.

const int N = 40;
int dp[N+5][10*N + 5]; // dp[i][j] = number of numbers of i digits with sum == j

void prep(){
	memset(dp,0,sizeof(dp));
	dp[0][0] = 1;
	
	for(int i = 1; i < N; i++){
		for(int j = 0; j < 10; j++){
			for(int s = j;s < 10*N;s++){
				dp[i][s] += dp[i-1][s-j];
			}
		}
	}
}


map<string,int> mp;
int g(string s){
	
	if(mp.find(s) != mp.end()){
		return mp[s];
	}
	
	int n = s.size();
	
	int &ans = mp[s];
	
	int sum = 0;
	
	for(int i = 0; i < n; i++){
		
		for(int j = (i == 0); j < s[i] - '0'; j++){
			int targ = 10 - (sum + j);
			if(targ >= 0){
				ans += dp[n-i-1][targ];
			}
		}
		sum += s[i] - '0';
	}
	
	ans += sum == 10;
	
	if(n > 1){
		ans += get(string(n-1,'9'));
	}
	
	return ans;
}

Introduction

Hello Codeforces,
I came across this problem a while back — https://cses.fi/problemset/task/1139/. Here the task is — "Given a colored, rooted tree, determine the number of distinct colors in each subtree of the tree". It is possible to solve this problem by flattening the tree and then using a binary indexed tree(now that subtrees have converted into subarrays). I will discuss an alternate technique to solve this, which can help in cases where it is useful to have the entire subtree as a set in a DFS call.

$$$O(n^2)$$$ solution

If $$$n$$$ were small, we could have stored for each node, the set of elements in its subtree. A DFS can be used to evaluate this, which will take $$$O(n^2log n)$$$ time and $$$O(n^2)$$$ space.

vector<int> color(n);
vector<set<int>> subtree(n);

void dfs(int i,int parent){
	subtree[i].insert(color[i]);
	for(int node : graph[i]){
		if(node != parent){
			dfs(node,i);
			subtree[i].insert(subtree[node].begin(),subtree[node].end());
		}
	}
}

The number of distinct colors in the subtree of the $$$i^{\text{th}}$$$ node is just the size of the set corresponding to it.

An Improvement

The problem with this technique is repetition. Each element occurs in $$$depth[i]$$$ different sets which is inefficient. For an improvement, we can consider the same idea used in heavy-light decomposition and union-find optimizations, which is "small-to-large merging".

We will use a DFS as before. To construct the set of all elements in the node's subtree, we will pick the child with the largest subtree set. Instead of making a new set for the current node, we will use the same set as the child with the largest subtree and insert all other children's sets in it. To avoid consuming extra space and time, we will use pointers to accomplish this task.

vector<int> color,distinct;
vector<set<int>*> subtree;
 
void dfs(int i,int parent = -1){
	int largest = -1;
	vector<int> children;
	for(int node : graph[i]){
		if(node != parent){
			dfs(node,i);
			children.push_back(node);
			if(largest == -1 || subtree[largest]->size() < subtree[node]->size()){
				largest = node;
			}
		}
	}
	if(largest == -1){
		subtree[i] = new set<int>; // new set for leaf node
	}
	else{
		subtree[i] = subtree[largest]; // largest sized child
	}
	
	for(int child : children){
		if(child == largest)continue;
		subtree[i]->insert(subtree[child]->begin(),subtree[child]->end());
	}
	subtree[i]->insert(color[i]);
	distinct[i] = subtree[i]->size();
}

Time complexity

Every time we copy an element over, the set it is now in will be at least 2 times larger than the set it was previously in. Hence the time complexity is $$$O(n log^2 n)$$$ (or $$$O(n log n)$$$ if we avoid using sets).

Thanks to -is-this-fft- for this

Conclusion

Here is my solution for the CSES task — https://cses.fi/paste/e75f5efc2e1cd2fc3c8a66/

This technique can be used in various other tree problems. Examples of such problems could be to find the most frequent element in each node's subtree or to find the pair with minimum difference in each node's subtree.