Codechef Devu and Locks

Devu and Locks

求有多少\(n\)位十进制数（可以有前导\(0\)），模\(p = 0\)，各个数位之和不超过\(m\)。

模\(998244353\)。

\(n ≤ 10^9, p ≤ 16, m ≤ 15000\)。

题解

任轩笛《杂题选讲》。

就是个倍增二维FFT。

\(f_{i,j,k}\)表示\(2^i\)位数，\(\bmod p=j\)，数位和是\(k\)的方案数。

\[f_{i,j_1,k_1}\times f_{i,j_2,k_2}\rightarrow f_{i+1,(j_1\times 10^{2^i}+j_2)\bmod p,k_1+k_2} \]

时间复杂度\(O(pm\log(pm)\log n)\)。

typedef vector<poly> poly2;

CO int N=1<<15;
int omg[2][N],rev[N];

void NTT(poly&F,int dir){
	int lim=F.size(),len=log2(lim);
	for(int i=0;i<lim;++i) rev[i]=rev[i>>1]>>1|(i&1)<<(len-1);
	for(int i=0;i<lim;++i)if(i<rev[i]) swap(F[i],F[rev[i]]);
	for(int i=1;i<lim;i<<=1)
		for(int j=0;j<lim;j+=i<<1)for(int k=0;k<i;++k){
			int t=mul(omg[dir][N/(i<<1)*k],F[j+i+k]);
			F[j+i+k]=add(F[j+k],mod-t),F[j+k]=add(F[j+k],t);
		}
	if(dir){
		int ilim=fpow(lim,mod-2);
		for(int i=0;i<lim;++i) F[i]=mul(F[i],ilim);
	}
}
poly operator+(CO poly&F,CO poly&G){
	int n=F.size()-1,m=G.size()-1;
	poly H(max(n,m)+1);
	for(int i=0;i<=max(n,m);++i)
		H[i]=add(i<=n?F[i]:0,i<=m?G[i]:0);
	return H;
}
poly2 reverse(CO poly2&F){
	int n1=F.size()-1,n2=F.front().size()-1;
	poly2 G(n2+1,poly(n1+1));
	for(int i=0;i<=n1;++i)for(int j=0;j<=n2;++j)
		G[j][i]=F[i][j];
	return G;
}
poly2 operator*(poly2 F,poly2 G){
	int n1=F.size()-1,n2=F.front().size()-1;
	int m1=G.size()-1,m2=G.front().size()-1;
	int lim1=1<<(int)ceil(log2(n1+m1+1)),lim2=1<<(int)ceil(log2(n2+m2+1));
	F.resize(lim1),G.resize(lim1);
	for(int i=0;i<lim1;++i) F[i].resize(lim2),G[i].resize(lim2);
	for(int i=0;i<lim1;++i) NTT(F[i],0),NTT(G[i],0);
	F=reverse(F),G=reverse(G);
	for(int i=0;i<lim2;++i) NTT(F[i],0),NTT(G[i],0);
	for(int i=0;i<lim2;++i)for(int j=0;j<lim1;++j)
		F[i][j]=mul(F[i][j],G[i][j]);
	for(int i=0;i<lim2;++i) NTT(F[i],1),NTT(G[i],1);
	F=reverse(F),G=reverse(G);
	for(int i=0;i<lim1;++i) NTT(F[i],1),NTT(G[i],1);
	F.resize(n1+m1+1);
	for(int i=0;i<=n1+m1;++i) F[i].resize(n2+m2+1);
	return F;
}

int main(){
	omg[0][0]=1,omg[0][1]=fpow(3,(mod-1)/N);
	omg[1][0]=1,omg[1][1]=fpow(omg[0][1],mod-2);
	for(int i=2;i<N;++i){
		omg[0][i]=mul(omg[0][i-1],omg[0][1]);
		omg[1][i]=mul(omg[1][i-1],omg[1][1]);
	}
	int n=read<int>(),P=read<int>(),m=read<int>();
	poly2 F(P,poly(m+1)),G(P,poly(m+1));
	F[0][0]=1;
	for(int i=0;i<=min(9,m);++i) G[i%P][i]=1;
	int power=10%P;
	for(;n;n>>=1){
		function<poly2(poly2,poly2)> merge=[&](CO poly2&F,CO poly2&G)->poly2{
			poly2 H(P,poly(m+1));
			for(int i=0;i<=P-1;++i) H[i*power%P]=H[i*power%P]+F[i];
			H=H*G;
			for(int i=0;i<=2*P-2;++i) H[i].resize(m+1);
			for(int i=P;i<=2*P-2;++i) H[i%P]=H[i%P]+H[i];
			H.resize(P);
			return H;
		};
		if(n&1) F=merge(F,G);
		G=merge(G,G);
		power=power*power%P;
	}
	for(int i=1;i<=m;++i) F[0][i]=add(F[0][i-1],F[0][i]);
	for(int i=0;i<=m;++i) printf("%d%c",F[0][i]," \n"[i==m]);
	return 0;
}

这题其实方法很多，下面介绍一种常数比较好的做法：

不同的位对数位和的贡献都是等价的，区别在于它们模\(p\)的值。然而实际上不等价的只有\(p\)类。

预处理出\(\text{num}_i\)表示模\(p = i\)的\(10^x\)种数。然后对每一种去做个倍增FFT，得到\(f_j\)表示\(\text{num}_i\)个\(0..9\)，和为\(j\)的方案数（这一步为\(O(pm \log m \log n)\)，\(\text{num}_i\)是都不同的，但是倍增过程可以共用以减小常数）。

第\(i\)类若和为\(j\)，对模\(p\)的贡献就是\(ij \bmod p\)。然后把\(p\)类用个二维FFT合起来。模\(p\)的那维很小，直接暴力卷积就行了。（要做\(p\)次二维FFT，每次DFT是\(O(m \log m ∗ p)\)的，卷积是\(O(mp^2)\)的）

总复杂度是\(O(pm \log m \log n + p^2m \log m + p^3m)\)。

CO int N=1<<15;
int omg[2][N],rev[N];

void NTT(poly&F,int dir){
	int lim=F.size(),len=log2(lim);
	for(int i=0;i<lim;++i) rev[i]=rev[i>>1]>>1|(i&1)<<(len-1);
	for(int i=0;i<lim;++i)if(i<rev[i]) swap(F[i],F[rev[i]]);
	for(int i=1;i<lim;i<<=1)
		for(int j=0;j<lim;j+=i<<1)for(int k=0;k<i;++k){
			int t=mul(omg[dir][N/(i<<1)*k],F[j+i+k]);
			F[j+i+k]=add(F[j+k],mod-t),F[j+k]=add(F[j+k],t);
		}
	if(dir){
		int ilim=fpow(lim,mod-2);
		for(int i=0;i<lim;++i) F[i]=mul(F[i],ilim);
	}
}
poly operator*(poly F,poly G){
	int n=F.size()+G.size()-1,lim=1<<int(ceil(log2(n)));
	F.resize(lim),NTT(F,0);
	G.resize(lim),NTT(G,0);
	for(int i=0;i<lim;++i) F[i]=mul(F[i],G[i]);
	NTT(F,1),F.resize(n);
	return F;
}

typedef array<array<int,100>,100> matrix;
int P,cnt[50];

matrix operator*(CO matrix&A,CO matrix&B){
	matrix C={};
	for(int k=0;k<=2*P-1;++k)
		for(int i=0;i<=2*P-1;++i)for(int j=0;j<=2*P-1;++j)
			C[i][j]=add(C[i][j],mul(A[i][k],B[k][j]));
	return C;
}
matrix pow(matrix B,int k){
	matrix A={};
	for(int i=0;i<=2*P-1;++i) A[i][i]=1;
	for(;k;k>>=1,B=B*B)
		if(k&1) A=A*B;
	return A;
}

poly B[30];

int main(){
	omg[0][0]=1,omg[0][1]=fpow(3,(mod-1)/N);
	omg[1][0]=1,omg[1][1]=fpow(omg[0][1],mod-2);
	for(int i=2;i<N;++i){
		omg[0][i]=mul(omg[0][i-1],omg[0][1]);
		omg[1][i]=mul(omg[1][i-1],omg[1][1]);
	}
	int n=read<int>();
	read(P);
	if(P==1) cnt[0]=n;
	else{
		matrix A={};
		for(int i=0;i<=P-1;++i) ++A[i][i*10%P],++A[i][i+P],++A[i+P][i+P];
		A=pow(A,n);
		for(int i=0;i<=P-1;++i) cnt[i]=A[1][i+P];
	}
	int m=read<int>();
	B[0].resize(m+1);
	for(int i=0;i<=min(9,m);++i) B[0][i]=1;
	for(int i=1;1<<i<=n;++i) B[i]=B[i-1]*B[i-1],B[i].resize(m+1);
	poly2 F(P,poly(m+1));
	F[0][0]=1;
	for(int i=0;i<=P-1;++i)if(cnt[i]){
		poly A(m+1);
		A[0]=1;
		for(int j=0;1<<j<=cnt[i];++j)if(cnt[i]>>j&1)
			A=A*B[j],A.resize(m+1);
		poly2 G(P,poly(m+1));
		for(int j=0;j<=m;++j) G[j*i%P][j]=A[j];
		int lim=1<<(int)ceil(log2(2*m+1));
		for(int j=0;j<=P-1;++j){
			F[j].resize(lim),NTT(F[j],0);
			G[j].resize(lim),NTT(G[j],0);
		}
		poly2 H(P,poly(lim));
		for(int j=0;j<=P-1;++j)for(int k=0;k<=P-1;++k)
			for(int l=0;l<lim;++l)
				H[(j+k)%P][l]=add(H[(j+k)%P][l],mul(F[j][l],G[k][l]));
		for(int j=0;j<=P-1;++j)
			NTT(H[j],1),H[j].resize(m+1),F[j]=H[j];
	}
	for(int i=1;i<=m;++i) F[0][i]=add(F[0][i-1],F[0][i]);
	for(int i=0;i<=m;++i) printf("%d%c",F[0][i]," \n"[i==m]);
	return 0;
}