CF1096G Lucky Tickets [FFT/NTT]

题意

一个$n$位数，每位可以是给出的$k$个数码中的一个数，可以有前导$0$,输出前$n/2$位之和与后$n/2$位之和相等的方案数，保证$n$是偶数。

输入的第一行是两个整数$n,k$，接下来的一行有$k$个数$d_1,d_2,\cdots,d_k(0\leq d_i\leq 9)$

输出一个数，为方案数模$998244353$的值。

---

求出用$n/2$位的数得到的和的方案数$f[n]$，答案为$$\sum_{i=1}^{maxN}f[i]^{2}$$

将数字$k$视为多项式的一项$x^{k}$，得到$$F(x)=\sum_{i=0}^{9}a_{i}x^{i}$$ 其中$a_{i}$是$0$或$1$。

那么一共可以选择$n/2$位，$F(x)$的$n/2$次就是答案，可以点值后快速幂再插值。

#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;
typedef long long ll;
const int N=8000010;
const ll P=998244353,g=3,gi=332748118;
int n,m,K,fn,r[N],l,x,mw;
ll ans,f[N],inv;
inline ll qpow(ll x,ll k) {ll r=1; for (; k; k>>=1ll,x=x*x%P) if (k&1ll) r=r*x%P; return r;}
inline void init(int n) {
    for (fn=1,l=0; fn<=n; fn<<=1) ++l;
    for (int i=0; i<fn; i++) r[i]=(r[i>>1]>>1)|((i&1)<<(l-1));
    inv=qpow((ll)fn,P-2ll);
}
void NTT(ll *c,ll t) {
    for (int i=0; i<fn; i++) if (i<r[i]) swap(c[i],c[r[i]]);
    for (int i=1; i<fn; i<<=1) {
        ll Wn=qpow(t,(P-1ll)/(ll)(i<<1));
        for (int j=0; j<fn; j+=(i<<1)) {
            ll w=1;
            for (int k=0; k<i; k++,w=w*Wn%P) {
                ll x=c[j+k],y=w*c[i+j+k]%P;
                c[j+k]=(x+y)%P; c[i+j+k]=(x+P-y)%P;
            }
        }
    }
}
int main() {
    scanf("%d%d",&n,&K); n>>=1;
    while (K--) scanf("%d",&x),f[x]=1ll;
    init(9*n);
    NTT(f,g);
    for (int i=0; i<fn; i++) f[i]=qpow(f[i],n);
    NTT(f,gi);
    for (int i=0; i<fn; i++) f[i]=f[i]*inv%P;
    for (int i=0; i<=fn; i++) ans=(ans+f[i]*f[i]%P)%P;
    printf("%lld\n",ans);
    return 0;
}