【bzoj4818】 Sdoi2017—序列計數

http://www.lydsy.com/JudgeOnline/problem.php?id=4818 (題目連接)

題意

　　一個長度爲$n$的序列，每一個元素是不超過$m$的正整數，且這$n$個數的和是$p$的倍數，這$n$個數中至少有一個是質數，問這樣的序列有多少個。

Solution

　　md嚇死我了，還覺得想錯了，$p^2\log n$的半天不敢寫=。=

　　$f[i][j]$表示忽略質數條件下的長度爲$i$，和$mod~p=j$的序列數；$g[i][j]$表示知足沒有一個數是質數的狀況下長度爲$i$，和$mod~p=j$的序列數。

　　而後這個東西倍增優化一下，每次$p^2$爆乘就行了。

細節

　　循環卷積。mdzz卡空間，明明原題空間是256M。。

代碼

#include<algorithm>
#include<iostream>
#include<cstdlib>
#include<cstring>
#include<cstdio>
#include<cmath>
#define LL long long
#define inf (1ll<<30)
#define Pi acos(-1.0)
#define free(a) freopen(a".in","r",stdin),freopen(a".out","w",stdout);
using namespace std;

const int maxn=10000010,maxm=1010,MOD=20170408;
int p[maxn];
bool vis[maxn<<1];
LL F[maxm],G[maxm],f[maxm],g[maxm],c[maxm];
int n,m,P;

void NTT(LL *a,LL *b,LL *r) {
	int N=P<<1;
	for (int i=0;i<N;i++) c[i]=0;
	for (int i=0;i<P;i++)
		for (int j=0;j<P;j++) (c[i+j]+=a[i]*b[j]%MOD)%=MOD;
	for (int i=0;i<P;i++) r[i]=(c[i]+c[i+P])%MOD;
}
int main() {
	scanf("%d%d%d",&n,&m,&P);
	vis[1]=1;
	for (int i=2;i<=m;i++) {
		if (!vis[i]) p[++p[0]]=i;
		for (int j=1;j<=p[0] && i*p[j]<=m;j++) {
			vis[i*p[j]]=1;
			if (i%p[j]==0) break;
		}
	}
	for (int i=1;i<=m;i++) {
		(++f[i%P])%=MOD;
		if (vis[i]) (++g[i%P])%=MOD;
	}
	F[0]=G[0]=1;
	for (;n;n>>=1) {
		if (n&1) NTT(F,f,F),NTT(G,g,G);
		NTT(f,f,f),NTT(g,g,g);
	}
	printf("%lld",(F[0]-G[0]+MOD)%MOD);
	return 0;
}