AGC009E Eternal Average

https://atcoder.jp/contests/agc009/tasks/agc009_e

将模型转化为一个 $k$ 叉树，每个点要么是叶子，要么恰好有 $k$ 个儿子
然后叶子节点有 $n+m$ 个，权值 0 或 1，非叶子节点的权值是所有子节点的平均值
这样根的权值就是最后剩的那个数
设根的深度为 $0$，权值为 1 的点深度分别是 $x_i$，为 0 的分别是 $y_i$
则权值为 0 的没贡献，根节点的权值是 $\sum k^{-x_i}$
如果 $n+m$ 个点都是权值为 $1$，则根权值也是 $1$，那么就推出了 $\sum k^{-x_i}+\sum k^{-y_i}=1$

此题即为求有多少个 $Z$，使得 $Z=\sum k^{-x_i},1-Z=\sum k^{-y_i}$
那么设 $k$ 进制小数 $Z=(0.z_1z_2\ldots z_l)_k$，因为 $Z$ 要写成 $\sum k^{-x_i}$，所以 $\sum z_i\le n$，也就是最多只有 $k$ 进制下 $n$ 个 $1$
那么从 $1$ 最多的情况往后推，每次进位，所有位的和减小 $k-1$，所以可以得出 $\sum z_i \equiv n \pmod{k-1}$

模拟一下小数减法的过程，再针对 $1-Z=\sum k^{-y_i}$ 这个式子，可以仿照上面的过程推出 $1+\sum(k-1-z_i)\le m$ 和 $1+\sum(k-1-z_i)\equiv m \pmod{k-1}$

所以就是要求有多少个小数 $Z$，使得满足一下几个式子

$$0\le z_i<k,i\in [1,l-1]$$

$$0<z_l<k$$

$$\sum z_i\le n$$

$$\sum z_i \equiv n \pmod{k-1}$$

$$1+\sum(k-1-z_i)\le m$$

$$1+\sum(k-1-z_i)\equiv m \pmod{k-1}$$

所以可以来一个 dp，用 $f(i,j,0/1)$ 表示考虑 $l=i$，也就是 $Z$ 有 $i$ 位时，$\sum z_i=j$，$z_i$ 等于/不等于 $0$ 的方案数
转移：$f(i,j,0)=f(i-1,j,0)+f(i-1,j,1)$
$f(i,j,1)=\sum_{x=1}^{k-1} f(i-1,j-x,0)+f(i-1,j-x,1)$

记一个前缀和优化一下就行了
因为没有规定 $l$ 的大小，所以每当枚举完一个 $i$，就去找所有符合上面几个式子要求的 $j$ 来加到答案里

另外一个细节，就是 $i$ 最大枚举到多少，其实通过联立第三个和第五个式子可以发现，$l\le \dfrac{n+m-1}{k-1}$，这也就是 $i$ 的最大值

#include<cstdio>
#include<algorithm>
#include<cmath>
#include<map>
#define reg register
inline int read(){
	register int x=0;register int y=0;
	register char c=getchar();
	while(c<'0'||c>'9'){if(c=='-')y=1;c=getchar();}
	while(c>='0'&&c<='9') x=x*10+c-48,c=getchar();
	return y?-x:x;
}
#define mod 1000000007
int n,m,k;
int sum[2006];
int f[4006][2006][2];
int main(){
	n=read();m=read();k=read();
	f[0][0][0]=1;
	int max=n+m,ans=0;
	for(reg int i=1;i<=max;i++){
		sum[0]=(f[i-1][0][0]+f[i-1][0][1])%mod;
		for(reg int j=1;j<=n;j++) sum[j]=((f[i-1][j][0]+f[i-1][j][1])%mod+sum[j-1])%mod;
		f[i][0][0]=sum[0];
		for(reg int j=1;j<=n;j++){
			f[i][j][0]=(sum[j]-sum[j-1]+mod)%mod;
			f[i][j][1]=(sum[j-1]-(j-k>=0?sum[j-k]:0)+mod)%mod;
		}
		for(reg int j=0;j<=n;j++)
			if(j%(k-1)==n%(k-1)&&1+(k-1)*i-j<=m&&(1+(k-1)*i-j)%(k-1)==m%(k-1))
				ans=(ans+f[i][j][1])%mod;
	}
	printf("%d",ans);
	return 0;
}