Codeforces632E Thief in a Shop（NTT + 快速幂）

题目

Source

http://codeforces.com/contest/632/problem/E

Description

A thief made his way to a shop.

As usual he has his lucky knapsack with him. The knapsack can contain k objects. There are n kinds of products in the shop and an infinite number of products of each kind. The cost of one product of kind i is ai.

The thief is greedy, so he will take exactly k products (it's possible for some kinds to take several products of that kind).

Find all the possible total costs of products the thief can nick into his knapsack.

Input

The first line contains two integers n and k (1 ≤ n, k ≤ 1000) — the number of kinds of products and the number of products the thief will take.

The second line contains n integers ai (1 ≤ ai ≤ 1000) — the costs of products for kinds from 1 to n.

Output

Print the only line with all the possible total costs of stolen products, separated by a space. The numbers should be printed in the ascending order.

Sample Input

3 2
1 2 3

5 5
1 1 1 1 1

3 3
3 5 11

Sample Output

2 3 4 5 6
5
9 11 13 15 17 19 21 25 27 33

分析

题目大概说给有n种价值各一的物品，每种数量都无限多，问取出k个物品能取出的物品价值和的所有情况。

用母函数解，价值为指数、存不存在为系数，构造多项式求k次幂即可。
这自然想到FFT+快速幂求，这样时间复杂度才够。

FFT直接求的话结果的系数最大到达1000¹⁰⁰⁰太爆炸了，当然也可以求一次卷积后非0指数重新赋值成1；不过我想着开头一次DFT结尾一次IDFT这样更快、更轻松点，所以用NTT了。。

我NTT模数取1004535809 WA在20，取998244353 WA在21。。看样子是系数取模后变为0了，数据叼叼的。。于是我就两个模数都取，然后4000多ms险过了。。

代码

#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
#define MAXN 1048576
 
//const long long P=50000000001507329LL; // 190734863287 * 2 ^ 18 + 1
long long P=1004535809; // 479 * 2 ^ 21 + 1
//const long long P=998244353; // 119 * 2 ^ 23 + 1
const int G=3;
 
long long mul(long long x,long long y){
	return (x*y-(long long)(x/(long double)P*y+1e-3)*P+P)%P;
}
long long qpow(long long x,long long k,long long p){
	long long ret=1;
	while(k){
		if(k&1) ret=mul(ret,x);
		k>>=1;
		x=mul(x,x);
	}
	return ret;
}
 
long long wn[25];
void getwn(){
	for(int i=1; i<=21; ++i){
		int t=1<<i;
		wn[i]=qpow(G,(P-1)/t,P);
	}
}
 
int len;
void NTT(long long y[],int op){
	for(int i=1,j=len>>1,k; i<len-1; ++i){
		if(i<j) swap(y[i],y[j]);
		k=len>>1;
		while(j>=k){
			j-=k;
			k>>=1;
		}
		if(j<k) j+=k;
	}
	int id=0;
	for(int h=2; h<=len; h<<=1) {
		++id;
		for(int i=0; i<len; i+=h){
			long long w=1;
			for(int j=i; j<i+(h>>1); ++j){
				long long u=y[j],t=mul(y[j+h/2],w);
				y[j]=u+t;
				if(y[j]>=P) y[j]-=P;
				y[j+h/2]=u-t+P;
				if(y[j+h/2]>=P) y[j+h/2]-=P;
				w=mul(w,wn[id]);
			}
		}
    }
    if(op==-1){
		for(int i=1; i<len/2; ++i) swap(y[i],y[len-i]);
		long long inv=qpow(len,P-2,P);
		for(int i=0; i<len; ++i) y[i]=mul(y[i],inv);
    }
}
void Convolution(long long A[],long long B[],int n){
	for(len=1; len<(n<<1); len<<=1);
	for(int i=n; i<len; ++i){
		A[i]=B[i]=0;
	}
	
	NTT(A,1); NTT(B,1);
	for(int i=0; i<len; ++i){
		A[i]=mul(A[i],B[i]);
	}
	NTT(A,-1);
}
 
long long A[MAXN],B[MAXN],C[MAXN];
long long cnt[MAXN];

int main(){
	getwn();
	int n,k,a;
	scanf("%d%d",&n,&k);
	int mx=0;
	for(int i=0; i<n; ++i){
		scanf("%d",&a);
		++cnt[a];
		mx=max(mx,a);
	}
	for(len=1; len<mx*k; len<<=1);
	
	memcpy(A,cnt,sizeof(cnt));
	NTT(A,1);
	memcpy(B,A,sizeof(B));
	--k;
	int tmp=k;
	while(k){
		if(k&1){
			for(int i=0; i<len; ++i) B[i]=mul(A[i],B[i]);
		}
		for(int i=0; i<len; ++i) A[i]=mul(A[i],A[i]);
		k>>=1;
	}
	NTT(B,-1);
	
	P=998244353;
	getwn();
	memcpy(A,cnt,sizeof(cnt));
	NTT(A,1);
	memcpy(C,A,sizeof(C));
	k=tmp;
	while(k){
		if(k&1){
			for(int i=0; i<len; ++i) C[i]=mul(A[i],C[i]);
		}
		for(int i=0; i<len; ++i) A[i]=mul(A[i],A[i]);
		k>>=1;
	}
	NTT(C,-1);
	
	for(int i=0; i<len; ++i){
		if(B[i] || C[i]) printf("%d ",i);
	}
	return 0;
}