【Codeforces 632E】Thief in a Shop
建生成函数,然后直接FFT+快速幂,中间不要算IDFT,最后再算。复杂度是\(O(WlogW)\)的,W表示最大值。取多个模数防止出现答案不为0但模数为0的情况。
#include <cmath> #include <cstdio> #include <cstring> #include <algorithm> using namespace std; typedef long long LL; const LL tp[3]={7340033,13631489,23068673},PN=1048576,tgn[3]={2187,1594323,7664329}; const int maxn=1050010; LL Pow(LL x,int pow,LL p) { LL res=1; for (;pow;pow>>=1,(x*=x)%=p) if (pow&1) (res*=x)%=p; return res; } int rev[maxn]; void FFT(int *a,int N,int f,int o) { LL p=tp[o],gn=tgn[o]; for (int i=0;i<N;i++) if (i<rev[i]) swap(a[i],a[rev[i]]); for (int i=1;i<N;i<<=1) { LL wn=Pow(gn,(f==1)?((PN/i)>>1):PN-((PN/i)>>1),p); for (int j=0;j<N;j+=(i<<1)) { LL w=1; for (int k=0;k<i;k++,(w*=wn)%=p) { int x=a[j+k],y=w*a[j+k+i]%p; a[j+k]=(x+y)%p,a[j+k+i]=(x-y+p)%p; } } } LL invN=Pow(N,p-2,p); if (f==-1) for (int i=0;i<N;i++) a[i]=invN*a[i]%p; } int x[3][maxn],xx[3][maxn]; int a[maxn]; int n,k,tempk,maxa; int main() { scanf("%d%d",&n,&k); for (int i=1;i<=n;i++) scanf("%d",&a[i]),maxa=max(maxa,a[i]); for (int i=1;i<=n;i++) xx[0][a[i]]++,xx[1][a[i]]++,xx[2][a[i]]++; int N=1;while (N<maxa*k+1) N<<=1; for (int i=0;i<N;i++) rev[i]=(rev[i>>1]>>1)|((i&1)*(N>>1)); x[0][0]=x[1][0]=x[2][0]=1; FFT(xx[0],N,1,0);FFT(xx[1],N,1,1);FFT(xx[2],N,1,2); FFT(x[0],N,1,0);FFT(x[1],N,1,1);FFT(x[2],N,1,2); tempk=k; while (tempk) { if (tempk&1) for (int i=0;i<N;i++) x[0][i]=(LL)x[0][i]*xx[0][i]%tp[0],x[1][i]=(LL)x[1][i]*xx[1][i]%tp[1],x[2][i]=(LL)x[2][i]*xx[2][i]%tp[2]; for (int i=0;i<N;i++) xx[0][i]=(LL)xx[0][i]*xx[0][i]%tp[0],xx[1][i]=(LL)xx[1][i]*xx[1][i]%tp[1],xx[2][i]=(LL)xx[2][i]*xx[2][i]%tp[2]; tempk>>=1; } FFT(x[0],N,-1,0);FFT(x[1],N,-1,1);FFT(x[2],N,-1,2); for (int i=0;i<=k*maxa;i++) if (x[0][i] || x[1][i] || x[2][i]) printf("%d ",i); return 0; }