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;
}