Luogu5339 [TJOI2019]唱、跳、rap和篮球 【生成函数，NTT】

当时看到这道题的时候我的脑子可能是这样的：

My left brain has nothing right, and my right brain has nothing left.

总之，看到"没有鸡你太美"这一类就直接想容斥，转化为”选出$i$个鸡你太美“

看到排列问题，直接想指数型生成函数。

设$m=\min(\frac{n}{4},a,b,c,d)$

我们使用万年不变的捆绑法，将鸡你太美当做整体考虑，即在$n-3i$个元素中选$i$个作为鸡你太美，再对其他四种进行全排列。

$$ans=\sum_{i=0}^m(-1)^i(n-4i)!\binom{n-3i}{i}\sum_{i_1\leq a-i}\sum_{i_2\leq b-i}\sum_{i_3\leq c-i}\sum_{i_4\leq d-i}\frac{1}{\prod i_j!}[\sum i_j=n-4i]$$

$$\sum_{i=0}^m(-1)^i\frac{(n-3i)!}{i!}\sum_{i_1\leq a-i}\sum_{i_2\leq b-i}\sum_{i_3\leq c-i}\sum_{i_4\leq d-i}\frac{1}{\prod i_j!}[\sum i_j=n-4i]$$

后面那一长串可以用NTT优化计算。

时间复杂度$O(n^2\log n)$，听说有直接dp的$O(n^2)$做法，但这个生成函数的做法应该是无脑多了。

#include<bits/stdc++.h>
#define Rint register int
using namespace std;
typedef long long LL;
const int N = 1 << 12, mod = 998244353, g = 3, gi = 332748118;
inline int kasumi(int a, int b){
    int res = 1;
    while(b){
        if(b & 1) res = (LL) res * a % mod;
        a = (LL) a * a % mod;
        b >>= 1;
    }
    return res;
}
int fac[N], inv[N];
inline void init(int n){
    fac[0] = 1;
    for(Rint i = 1;i <= n;i ++) fac[i] = (LL) i * fac[i - 1] % mod;
    inv[n] = kasumi(fac[n], mod - 2);
    for(Rint i = n;i;i --) inv[i - 1] = (LL) inv[i] * i % mod;
}
int rev[N];
inline int calrev(int n){
    int limit = 1, L = -1;
    while(limit <= n){limit <<= 1; L ++;}
    for(Rint i = 0;i < limit;i ++) rev[i] = (rev[i >> 1] >> 1) | ((i & 1) << L);
    return limit;
}
inline void NTT(int *A, int limit, int type){
    for(Rint i = 0;i < limit;i ++) if(i < rev[i]) swap(A[i], A[rev[i]]);
    for(Rint mid = 1;mid < limit;mid <<= 1){
        int Wn = kasumi(type == 1 ? g : gi, (mod - 1) / (mid << 1));
        for(Rint j = 0;j < limit;j += mid << 1){
            int w = 1;
            for(Rint k = 0;k < mid;k ++, w = (LL) w * Wn % mod){
                int x = A[j + k], y = (LL) w * A[j + k + mid] % mod;
                A[j + k] = (x + y) % mod;
                A[j + k + mid] = (x - y + mod) % mod;
            }
        }
    }
    if(type == -1){
        int inv = kasumi(limit, mod - 2);
        for(Rint i = 0;i < limit;i ++)
            A[i] = (LL) A[i] * inv % mod;
    }
}
int n, a, b, c, d, m, ans, A[N], B[N], C[N], D[N];
int main(){
    scanf("%d%d%d%d%d", &n, &a, &b, &c, &d);
    init(a + b + c + d);
    m = min(n >> 2, min(min(a, b), min(c, d)));
    for(Rint i = 0;i <= m;i ++){
        int limit = calrev(a + b + c + d - (i << 2));
        for(Rint j = 0;j < limit;j ++){
            A[j] = inv[j] * (j <= a - i);
            B[j] = inv[j] * (j <= b - i);
            C[j] = inv[j] * (j <= c - i);
            D[j] = inv[j] * (j <= d - i);
        }
        NTT(A, limit, 1); NTT(B, limit, 1); NTT(C, limit, 1); NTT(D, limit, 1);
        for(Rint j = 0;j < limit;j ++) A[j] = (LL) A[j] * B[j] % mod * C[j] % mod * D[j] % mod;
        NTT(A, limit, -1);
        int tmp = (LL) A[n - (i << 2)] * fac[n - 3 * i] % mod * inv[i] % mod;
        if(i & 1) ans = (ans - tmp + mod) % mod;
        else ans = (ans + tmp) % mod;
    }
    printf("%d", ans);
}