CF632E Thief in a Shop [ 生成函数 + NTT ]

传送门

生成函数好巧妙啊, 如果能买到, 系数为1, 否则为0

所以最开始的生成函数就是 , 然后自己跟自己卷k次就可以了

NTT 用 998244353 和 1004535809 做两次, 就可以避免冲突

#include<bits/stdc++.h>
#define N 1050000
using namespace std;
typedef long long ll;
int n, k, vis[N]; ll A[N], B[N];
int up, bit, rev[N];
const int Mod1 = 998244353, Mod2 = 1004535809, G = 3;
ll add(ll a, ll b, ll Mod){ return (a+b) % Mod;}
ll mul(ll a, ll b, ll Mod){ return (a*b) % Mod;}
void Init(int len){ up = 1, bit = 0;
	while(up != len) up <<= 1, bit++;
	for(int i=0; i<up; i++) rev[i] = (rev[i>>1]>>1) | ((i&1) << (bit-1));
}
ll power(ll a, ll b, ll Mod){
	ll ans = 1; for(;b;b>>=1){
		if(b&1) ans = mul(ans, a, Mod); a = mul(a, a, Mod);
	} return ans;
}
void NTT(ll *a, int Mod, int flag){
	for(int i=0; i<up; i++) if(i < rev[i]) swap(a[i], a[rev[i]]);
	for(int i=1; i<up; i<<=1){
		ll wn = power(G, (Mod-1)/(i<<1), Mod);
		if(flag == -1) wn = power(wn, Mod-2, Mod);
		for(int j=0; j<up; j+=(i<<1)){
			ll w = 1;
			for(int k=0; k<i; k++, w=mul(w, wn, Mod)){
				ll x = a[k+j], y = mul(w, a[k+j+i], Mod);
				a[k+j] = add(x, y, Mod); a[k+j+i] = add(x, Mod-y, Mod);
			}
		}
	}
	if(flag == -1){
		ll inv = power(up, Mod-2, Mod);
		for(int i=0; i<up; i++) a[i] = mul(a[i], inv, Mod);
	}
}
int main(){
	scanf("%d%d", &n, &k);
	for(int i=1; i<=n; i++){
		int x; scanf("%d", &x); A[x] = B[x] = 1;
	} Init(1 << 20); 
	NTT(A, Mod1, 1); for(int i=0; i<up; i++) A[i] = power(A[i], k, Mod1);
	NTT(A, Mod1, -1); for(int i=0; i<up; i++) if(A[i]) vis[i] = 1;
	NTT(B, Mod2, 1); for(int i=0; i<up; i++) B[i] = power(B[i], k, Mod2);
	NTT(B, Mod2, -1); for(int i=0; i<up; i++) if(B[i]) vis[i] = 1;
	for(int i=1; i<up; i++) if(vis[i]) printf("%d ", i);
	return 0; 
}