Codeforces - 1096G - Lucky Tickets - NTT

https://codeforc.es/contest/1096/problem/G

把数组分成前后两半，那么前半部分的各个值的表示方案的平方的和就是答案。

这些数组好像可以dp出来。

一开始设dp[i]数组表示1<<i位的各个值，那么做16次NTT然后把n分解下去就求出来答案了。总共要30多次convolution，比较大的NTT少说有6次。

后来发现dp数组可以只记录一次，用完就可以接着用。

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

const int MAXN = 5e6, mod = 998244353;

inline int pow_mod(ll x, int n) {
    ll res;
    for(res = 1; n; n >>= 1, x = x * x % mod)
        if(n & 1)
            res = res * x % mod;
    return res;
}

inline int add_mod(int x, int y) {
    x += y;
    return x >= mod ? x - mod : x;
}

inline int sub_mod(int x, int y) {
    x -= y;
    return x < 0 ? x + mod : x;
}

void NTT(int a[], int n, int op) {
    for(int i = 1, j = n >> 1; i < n - 1; ++i) {
        if(i < j)
            swap(a[i], a[j]);
        int k = n >> 1;
        while(k <= j) {
            j -= k;
            k >>= 1;
        }
        j += k;
    }
    for(int len = 2; len <= n; len <<= 1) {
        int g = pow_mod(3, (mod - 1) / len);
        for(int i = 0; i < n; i += len) {
            int w = 1;
            for(int j = i; j < i + (len >> 1); ++j) {
                int u = a[j], t = 1ll * a[j + (len >> 1)] * w % mod;
                a[j] = add_mod(u, t), a[j + (len >> 1)] = sub_mod(u, t);
                w = 1ll * w * g % mod;
            }
        }
    }
    if(op == -1) {
        reverse(a + 1, a + n);
        int inv = pow_mod(n, mod - 2);
        for(int i = 0; i < n; ++i)
            a[i] = 1ll * a[i] * inv % mod;
    }
}

int A[MAXN + 5], B[MAXN + 5];

int pow2(int x) {
    int res = 1;
    while(res < x)
        res <<= 1;
    return res;
}

void convolution(int a[], int b[], int asize, int bsize, int c[], int &csize) {
    int n = pow2(asize + bsize - 1);
    for(int i = 0; i < n; ++i) {
        A[i] = i < asize ? a[i] : 0;
        B[i] = i < bsize ? b[i] : 0;
    }
    NTT(A, n, 1);
    NTT(B, n, 1);
    for(int i = 0; i < n; ++i)
        A[i] = 1ll * A[i] * B[i] % mod;
    NTT(A, n, -1);
    csize = n;
    for(int i = 0; i < n; ++i)
        c[i] = A[i];
    return;
}

int dp[MAXN], dpsize;   //dp[i]：1<<i位能表示的各个位数
int ans[MAXN], anssize;

int main() {
#ifdef Yinku
    freopen("Yinku.in", "r", stdin);
#endif // Yinku
    cout<<(1<<20)<<endl;
    int n, k;
    scanf("%d%d", &n, &k);
    for(int i = 0; i < k; ++i) {
        int tmp;
        scanf("%d", &tmp);
        dp[tmp] = 1;
    }
    dpsize = 10;
    int n2 = n >> 1;
    bool fi = true;
    while(n2) {
        if(n2 & 1) {
            if(!fi)
                convolution(ans, dp, anssize, dpsize, ans, anssize);
            else {
                fi = false;
                for(int j = 0; j < dpsize; j++)
                    ans[j] = dp[j];
                anssize = dpsize;
            }
        }
        n2 >>= 1;
        if(n2)
            convolution(dp, dp, dpsize, dpsize, dp, dpsize);
    }
    ll res = 0;
    for(int i = 0; i < anssize; i++) {
        res += 1ll * ans[i] * ans[i] % mod;
    }
    printf("%lld\n", res % mod);
    return 0;
}

但其实不需要用到convolution？？？

注意到convolution的本质其实是用NTT变成点值，然后用点值加法再用NTT变回来。

其实注意到

void convolution(int a[], int b[], int asize, int bsize, int c[], int &csize) {
    int n = pow2(asize + bsize - 1);
    for(int i = 0; i < n; ++i) {
        A[i] = i < asize ? a[i] : 0;
        B[i] = i < bsize ? b[i] : 0;
    }
    NTT(A, n, 1);
    NTT(B, n, 1);
    for(int i = 0; i < n; ++i)
        A[i] = 1ll * A[i] * B[i] % mod;
    NTT(A, n, -1);
    csize = n;
    for(int i = 0; i < n; ++i)
        c[i] = A[i];
    return;
}

里面，两个多项式卷积实际上只是对应位置做乘法。

那么只要把对应位置的乘法一次全部做完就可以了。

从另一个角度想，要是把这个数组本身视作一个多项式（生成函数），\(a_ix^i\)中，\(a_i\)就是和为\(i\)的方案数。那么

\(F(x)=\sum\limits_{i=0}^{9}a_ix^i\) 就是1位数的选法，选n次那就是\(F^n(x)\)

这个就直接点值化之后对点值直接快速幂然后再插值 。

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

const int MAXN = 1<<21, mod = 998244353;

inline int pow_mod(ll x, int n) {
    ll res;
    for(res = 1; n; n >>= 1, x = x * x % mod)
        if(n & 1)
            res = res * x % mod;
    return res;
}

inline int add_mod(int x, int y) {
    x += y;
    return x >= mod ? x - mod : x;
}

inline int sub_mod(int x, int y) {
    x -= y;
    return x < 0 ? x + mod : x;
}

void NTT(int a[], int n, int op) {
    for(int i = 1, j = n >> 1; i < n - 1; ++i) {
        if(i < j)
            swap(a[i], a[j]);
        int k = n >> 1;
        while(k <= j) {
            j -= k;
            k >>= 1;
        }
        j += k;
    }
    for(int len = 2; len <= n; len <<= 1) {
        int g = pow_mod(3, (mod - 1) / len);
        for(int i = 0; i < n; i += len) {
            int w = 1;
            for(int j = i; j < i + (len >> 1); ++j) {
                int u = a[j], t = 1ll * a[j + (len >> 1)] * w % mod;
                a[j] = add_mod(u, t), a[j + (len >> 1)] = sub_mod(u, t);
                w = 1ll * w * g % mod;
            }
        }
    }
    if(op == -1) {
        reverse(a + 1, a + n);
        int inv = pow_mod(n, mod - 2);
        for(int i = 0; i < n; ++i)
            a[i] = 1ll * a[i] * inv % mod;
    }
}

int pow2(int x) {
    int res = 1;
    while(res < x)
        res <<= 1;
    return res;
}

int A[MAXN + 5];

int main() {
#ifdef Yinku
    freopen("Yinku.in", "r", stdin);
#endif // Yinku
    int n, k;
    scanf("%d%d", &n, &k);
    for(int i = 0; i < k; ++i) {
        int tmp;
        scanf("%d", &tmp);
        A[tmp] = 1;
    }
    n>>=1;
    int maxn=pow2(n*9);
    NTT(A,maxn,1);
    for(int i = 0; i < maxn; i++)
        A[i]=pow_mod(A[i],n);
    NTT(A,maxn,-1);
    ll res=0;
    for(int i = 0; i < maxn; i++)
        res += 1ll * A[i] * A[i] % mod;
    printf("%lld\n", res % mod);
    return 0;
}