CF1096G Lucky Tickets [生成函数+NTT]

传送门

我们如果能求出 n/2 个数能组合出的某个数的方案数 v, 那么v^2 就是这个数的贡献

发现可以用生成函数, 类似下面这样构造

组合成9的方案数就是2, 然后NTT就可以了

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