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;

}