Red-White Fence（计数问题，FFT/NTT)

描述

一个围栏由一个红板和几块（可能为0）块白板组成。红板是最长的一块，并且左边的白板长度严格递增，右边严格递减。
这些板子从高度为0，从左到右连续紧密放置，组成一个多边形。
给出询问q，每个给出参数Q，满足多边形周长是Q的板子放置的方案数（相同位置长度相同视作一种方案）。

链接

---

分析

首先来考虑一下围栏的周长是多少，假设围栏由 $m$ 块白板和 1 块长度为 $L$ 的红板组成，那么很容易发现这样的围栏周长 $P=2\ast (L+m+1)$​。

现在来考虑用一块长度为 $L$​ 的红板，能够组成周长为 $P$​ 的围栏的方案数有多少。由之前推得得式子，要组成这样得围栏，需要的白板个数就是 $\frac{P}{2}-L-1$​。

因为红板得数量最多只有 5，所以迭代枚举红板的长度 $L$，并且对每个询问的周长 $Q_i$，计算用长度为 $L$ 的红板与 $\frac{Q_i}{2}-L-1$​ 块白板组成围栏的方案数即可。

现在的问题是，在长度小于红板长度 $L$ 的白板中，用 $k$​ 块与红板组成围栏得方案数是多少？

首先假设有 $n$​ 块长度小于 $L$ 的白板，并且每种长度只有一个白板，即 $n$​ 块白板长度互不相同。这种情况，首先是要从 $n$ 个白板中选 $k$ 个出来：$C_n^k$，然后对于每块白板，它可以选择放在红板左边或者右边: $2^k$。因此方案数是 $C_n^k{2}^k$​。

然后再来考虑一下相反情况，在长度小于红板长度的白板中，对于这些白板中的每种长度，都至少有两块白板是这种长度的。假设有 $m$​ 种不同长度的白板小于红板，每种长度白板，我们其实是有四种选择，1.不选这种长度的白板，2.选1个这种长度的白板并放在红板左侧，3.选1个这种长度的白板并放在红板右侧，4. 选 2 个这种长度的白板分别放在红板左侧右侧。因此可以这样考虑，把 $m$​​​​​ 种白板中的每一种白板拆成 2 个点，分别表示选这种白板并放在左侧，与选这种白板并放在右侧。那这种情况最后的方案数就是$C_{2m}^k$​（挺巧妙的）。

接下来就是综合考虑上面的两种情况了，既有长度唯一的白板，又有长度不唯一的白板。那么考虑枚举选取几个长度唯一的白板就容易得出答案为 ${\sum }_{i=0}^k{C}_{2m}^{k-i}{C}_n^i{2}^i$​。

至此，按照这个过程求解的时间复杂度是 $O(kqn)$​​，是会超时的。

接下来就是类似于整数拆分的技巧，用多项式优化。

对两个多项式 ${\sum }_{i=0}^n{C}_n^i{2}^i{x}^i$​​ 和 ${\sum }_{i=0}^{2m}{C}_{2m}^i{x}^i$​​​ 卷积一下即可，选 $k$​​ 个白板的答案就是结果多项式中 $x_k$​​​​ 的系数。用 $FFT$​​ 或者 $NTT$​​ 都可以求卷积，但是本题的模数特殊，用 $NTT$​​​ 更适合一些。

最终时间复杂度是 $O(k(nlogn+q))$​。​

( 注意：时间复杂度中的 $k,q,n$​​ 是题面中的变量，并非题解中指代的变量 )

代码

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

typedef long long ll;

const int N = 300010;
const int mod = 998244353;
const int G = 3;

int a[N], b[N], Q[N], ans[N], fact[N], rfact[N];
int p1[N << 2], p2[N << 2], rev[N << 2];

int qmi(int a, int b, int p)
{
    int ans = 1;
    for (; b; b >>= 1)
    {
        if (b & 1) ans = (ll)ans * a % mod;
        a = (ll)a * a % mod;
    }
    return ans;
}

int C(int n, int m)
{
    if (n < m) return 0;
    return (ll)fact[n] * rfact[m] % mod * rfact[n - m] % mod;
}

void getpoly(int a[], int n, int r)
{
    for (int i = 0; i <= n; i++)
        a[i] = (ll)C(n, i) * qmi(r, i, mod) % mod;
}

void ntt(int a[], int tot, int sgn, int p)
{
    for (int i = 0; i < tot; i++)
        if (i < rev[i])
            swap(a[i], a[rev[i]]);
    int g = sgn == 1 ? G : qmi(G, p - 2, p);
    for (int mid = 1, t = 1; mid < tot; mid <<= 1, t++)
    {
        int wn = qmi(g, (p - 1) >> t, p);
        for (int i = 0; i < tot; i += mid << 1)
        {
            int w = 1;
            for (int j = 0; j < mid; j++, w = (ll)w * wn % p)
            {
                int x = a[i + j], y = (ll)w * a[i + mid + j] % p;
                a[i + j] = (x + y) % p;
                a[i + mid + j] = ((x - y) % p + p) % p;
            }
        }
    }
    if (sgn == -1)
    {
        int invtot = qmi(tot, p - 2, p);
        for (int i = 0; i < tot; i++)
            a[i] = (ll)a[i] * invtot % p;
    }
}

void mul_poly(int a[], int n, int b[], int m)
{
    int tot = 1, bit = 0;
    while (tot < n + m + 1) tot <<= 1, ++bit;
    for (int i = 0; i < tot; i++)
        rev[i] = (rev[i >> 1] >> 1) | ((i & 1) << (bit - 1));

    ntt(a, tot, 1, mod); ntt(b, tot, 1, mod);
    for (int i = 0; i < tot; i++) a[i] = (ll)a[i] * b[i] % mod;
    ntt(a, tot, -1, mod);
}

int main()
{
    fact[0] = fact[1] = rfact[0] = rfact[1] = 1;
    for (int i = 2; i < N; i++) 
    {
        fact[i] = (ll)fact[i - 1] * i % mod;
        rfact[i] = (ll)(mod - mod / i) * rfact[mod % i] % mod;
    }
    for (int i = 2; i < N; i++) rfact[i] = (ll)rfact[i - 1] * rfact[i] % mod;

    int n, k;
    scanf("%d%d", &n, &k);
    map<int, int> mp;
    for (int i = 1; i <= n; i++) 
    {   
        scanf("%d", &a[i]);
        mp[a[i]]++;
    }
    for (int i = 1; i <= k; i++) scanf("%d", &b[i]);
    int q;
    scanf("%d", &q);
    for (int i = 1; i <= q; i++) scanf("%d", &Q[i]);

    for (int i = 1; i <= k; i++)
    {
        int red = b[i];
        int cnt1 = 0, cnt2 = 0;
        for (auto x : mp)
        { 
            if (x.first >= red) break;
            if (x.second == 1) cnt1++;
            else cnt2++;
        }
        memset(p1, 0, sizeof(p1));
        memset(p2, 0, sizeof(p2));
        getpoly(p1, cnt1, 2);
        getpoly(p2, 2 * cnt2, 1);
        mul_poly(p1, cnt1, p2, 2 * cnt2);

        for (int j = 1; j <= q; j++)
        {
            int num = Q[j] / 2 - red - 1;
            if (num >= 0 && num <= cnt1 + 2 * cnt2)
                ans[j] = (ans[j] + p1[num]) % mod;
        }
    }

    for (int i = 1; i <= q; i++)
        printf("%d\n", ans[i]);

    return 0;
}

因为自己也是初学FFT/NTT，在写代码时也遇到了一些的问题，诸如两数相乘忘记开long long，做NTT时误把p1数组写成了a数组，白白浪费了很多时间，所以还需继续努力！