codeforces 632E. Thief in a Shop(生成函数)

传送门
题意简述:给 n n n个物件,物件 i i i有一个权值 a i a_i ai,可以选任意多个。现在要求选出 k k k个物件出来(允许重复)问最后得到的权值和的种类数。
n , k , a i ≤ 1000 n,k,a_i\le1000 n,k,ai1000


思路:
这是一道很显然的生成函数,我们把选一个物件的生成函数给列出来,然后取它的 k k k次方就是答案。
显然可以上一波 f f t fft fft 成功T飞
在博主卡场无果之后换成了 n t t ntt ntt,然后发现 c f cf cf的强力数据同时卡掉了 998244353 998244353 998244353 1004535809 1004535809 1004535809这两个模数,于是菜鸡博主兴奋地跑了两次 n t t ntt ntt取并集,并且兴奋地发现自己过了。
代码:

#include
#define ri register int
using namespace std;
typedef long long ll;
int tim,lim,mod;
vector<int>pos,A,B;
inline void init(const int&up){
	tim=0,lim=1;
	while(lim<=up)lim<<=1,++tim;
	pos.resize(lim),A.resize(lim),B.resize(lim);
	for(ri i=0;i<lim;++i)pos[i]=(pos[i>>1]>>1)|((i&1)<<(tim-1));
}
inline int add(int a,int b){return a+b>=mod?a+b-mod:a+b;}
inline int dec(int a,int b){return a>=b?a-b:a-b+mod;}
inline int mul(int a,int b){return (ll)a*b%mod;}
inline int ksm(int a,int p){int ret=1;for(;p;p>>=1,a=mul(a,a))if(p&1)ret=mul(ret,a);return ret;}
inline void ntt(vector<int>&a,const int&type){
	for(ri i=0;i<lim;++i)if(i<pos[i])swap(a[i],a[pos[i]]);
	int wn,w,a0,a1;
	for(ri mid=1,typ=type==1?3:(mod+1)/3,mult=(mod-1)/2;mid<lim;mid<<=1,mult>>=1){
		wn=ksm(typ,mult);
		for(ri j=0,len=mid<<1;j<lim;j+=len){
			w=1;
			for(ri k=0;k<mid;++k,w=mul(w,wn)){
				a0=a[j+k],a1=mul(w,a[j+k+mid]);
				a[j+k]=add(a0,a1),a[j+k+mid]=dec(a0,a1);
			}
		}
	}
	if(type==-1)for(ri i=0,inv=ksm(lim,mod-2);i<lim;++i)a[i]=mul(a[i],inv);
}
const int N=1005;
int n,k,mx=0;
struct poly{
	vector<int>a;
	poly(int k=0,int w=0){a.resize(k+1),a[k]=w;}
	inline int&operator[](const int&k){return a[k];}
	inline const int&operator[](const int&k)const{return a[k];}
	inline int deg()const{return a.size()-1;}
	inline poly extend(const int&k){poly ret=*this;return ret.a.resize(k+1),ret;}
	friend inline poly operator^(const poly&a,const int&k){
		init(mx+1);
		int p=k-1;
		poly ret=a;
		for(ri i=0;i<=a.deg();++i)A[i]=a[i],B[i]=ret[i];
		for(ri i=a.deg()+1;i<lim;++i)A[i]=B[i]=0;
		ntt(A,1),ntt(B,1);
		while(p){
			if(p&1)for(ri i=0;i<lim;++i)B[i]=mul(B[i],A[i]);
			for(ri i=0;i<lim;++i)A[i]=mul(A[i],A[i]);
			p>>=1;
		}
		return ntt(B,-1),ret.a=B,ret.extend(mx);
	}
};
inline int read(){
	int ans=0;
	char ch=getchar();
	while(!isdigit(ch))ch=getchar();
	while(isdigit(ch))ans=(ans<<3)+(ans<<1)+(ch^48),ch=getchar();
	return ans;
}
int main(){
	n=read(),k=read();
	poly a(1000);
	for(ri v,i=1;i<=n;++i)v=read(),mx=max(mx,v),++a[v];
	a=a.extend(mx),mx*=k;
	poly res1,res2;
	mod=998244353,res1=(a^k);
	mod=1004535809,res2=(a^k);
	for(ri i=0;i<=mx;++i)if(res1[i]||res2[i])cout<<i<<' ';
	return 0;
}

你可能感兴趣的:(数学,#,快速幂,#,ntt,#,生成函数)