题解 P5339 【[TJOI2019]唱、跳、rap和篮球】

题目链接

Solution [TJOI2019]唱、跳、rap和篮球

题目大意：有 (a) 个人喜欢唱，(b) 个人喜欢跳，(c) 个人喜欢 rap，(d) 个人喜欢篮球。你要从中选出 (n) 个人组成一个排列，如果存在一个位置 (i)，使得 (i,i+1,i+2,i+3) 位置的人依次喜欢唱、跳、rap、篮球，队伍就会不和谐。求和谐排列数，相同喜好的学生之间不可区分。

二项式反演、计数

---

分析：

设 (f(i)) 为恰好有 (i) 组人讨论 cxk 的方案数，(U) 为所有方案数。方便起见，记 (lfloor frac{n}{4} floor =lim)。

那么 (ans=U-sum_{i=1}^{lim}f(i))

直接计算 (f) 比较困难，我们考虑二项式反演。

令 (g(i)) 表示我们钦定 (i) 个位置，以这些位置开始的若干 (4) 人组都讨论 cxk 的方案数。

可以改写答案为 (ans=g(0)-sum_{i=1}^{lim}f(i))

那么有 (g(x)=sum_{i=x}^{lim}inom{i}{x}f(i))
。对于每个恰好为 (i) 的方案，我们有 (inom{i}{x}) 种方法从里面钦定。

反演可知 (f(x)=sum_{i=x}^{lim}(-1)^{i-x}inom{i}{x}g(i))

考虑如何算 (g)。

显然讨论 cxk 的 (4) 人组是互不相交的，因此我们可以将他们捆绑在一起，视作整体。

也就是 (g(x))，有 (n-4x+x=n-3x) 个位置，这个时候钦定 (x) 个位置的方案为 (inom{n-3x}{x})。剩下的 (n-4x) 个人随意排列，问题和计算 (g(0)) 没有区别。

而 (g(0)) 是非常好计算的，假设只有喜欢唱、跳的人

从喜欢唱的人里面选 (x) 个的方案为 (A[x])，从喜欢跳的人里面选 (x) 个的方案为 (B[x])。

那么从唱跳里面选 (n) 个人组成合法排列的方案为 (g[n]=sum_{i=0}^ninom{n}{i}A[i]B[n-i])

这个方法很容易推广到四种，也就是将它们的 ( ext{EGF}) 卷起来就行了。

求 (g(x)) 是 (nlogn) 的，方便起见反演求 (f) 直接 (n^2) 算，所以总的复杂度是 (O(n^2logn)) 的

#include <cstdio>
#include <cctype>
#include <algorithm>
#include <vector>
#define debug(...) fprintf(stderr,__VA_ARGS__)
#pragma GCC optmize(2)
using namespace std;
typedef long long ll;
constexpr int maxn = 5e3,maxm = 1024,mod = 998244353,G = 3,invG = 332748118,inf = 0x7fffffff;
constexpr inline int add(const int a,const int b){return (a + b) % mod;}
constexpr inline int sub(const int a,const int b){return (a - b + mod) % mod;}
constexpr inline int mul(const int a,const int b){return (1ll * a * b) % mod;}
constexpr inline int qpow(int base,int b){//-std=c++17
	int res = 1;
	while(b){
		if(b & 1)res = mul(res,base);
		base = mul(base,base);
		b >>= 1;
	}
	return res;
}
constexpr inline int inv(const int x){return qpow(x,mod - 2);}//-std=c++17
constexpr inline int calc(const int a,const int b){return mul(a,inv(b));}

int tr[maxn << 1],len;
void gettr(const int x){
	for(len = 1;len < x;len <<= 1);
	for(int i = 0;i < len;i++)tr[i] = (tr[i >> 1] >> 1) | ((i & 1) ? (len >> 1) : 0);
}
struct poly : std::vector<int>{
	#define f (*this)
	using std::vector<int>::vector;
	void ntt(const int flg = 1){
		const int n = size();
		for(int i = 0;i < n;i++)
			if(i < tr[i])std::swap(f[i],f[tr[i]]);
		for(int p = 2;p <= n;p <<= 1){
			const int unit = qpow(flg == 1 ? G : invG,(mod - 1) / p);
			const int len = p >> 1;
			for(int k = 0;k < n;k += p){
				int now = 1;
				for(int l = k;l < k + len;l++){
					const int tt = mul(f[l + len],now);
					f[l + len] = sub(f[l],tt);
					f[l] = add(f[l],tt);
					now = mul(now,unit);
				}
			}
		}
		if(flg == -1){
			const int inv = ::inv(n);
			for(int i = 0;i < n;i++)f[i] = mul(f[i],inv);
		}
	}
	poly operator * (const poly &g)const{
		poly res;res.resize(size());
		for(unsigned int i = 0;i < size();i++)res[i] = mul(f[i],g[i]);
		return res;
	}
	#undef f
};
int binom[maxm][maxm],fac[maxm],facinv[maxm],lim,n,a,b,c,d,mi;
void init(){
	binom[0][0] = 1;
	for(int i = 1;i < maxm;i++){
		binom[i][0] = 1;
		for(int j = 1;j <= i;j++)binom[i][j] = add(binom[i - 1][j - 1],binom[i - 1][j]);
	}
	fac[0] = 1;
	for(int i = 1;i < maxm;i++)fac[i] = mul(fac[i - 1],i);
	facinv[maxm - 1] = inv(fac[maxm - 1]);
	for(int i = maxm - 2;i >= 0;i--)facinv[i] = mul(facinv[i + 1],i + 1);
}
int sgn(const int x){return (x & 1) ? mod - 1 : 1;}
int g(const int x){
	if(x > (n / 4))return 0;
	if(x > mi)return 0;
	poly a,b,c,d;
	a.resize(::a + 1);
	b.resize(::b + 1);
	c.resize(::c + 1);
	d.resize(::d + 1);
	for(int i = 0;i <= ::a - x;i++)a[i] = mul(1,facinv[i]);
	for(int i = 0;i <= ::b - x;i++)b[i] = mul(1,facinv[i]);
	for(int i = 0;i <= ::c - x;i++)c[i] = mul(1,facinv[i]);
	for(int i = 0;i <= ::d - x;i++)d[i] = mul(1,facinv[i]);
	gettr(a.size() + b.size() + c.size() + d.size() - 3);
	a.resize(len),b.resize(len),c.resize(len),d.resize(len);
	a.ntt();b.ntt();c.ntt();d.ntt();
	poly res = a * b * c * d;
	res.ntt(-1);
	return mul(mul(res[n - 4 * x],fac[n - 4 * x]),binom[n - 3 * x][x]);
}
int main(){
#ifndef ONLINE_JUDGE
	freopen("fafa.in","r",stdin);
#endif
	init();

	scanf("%d %d %d %d %d",&n,&a,&b,&c,&d);
	mi = min({a,b,c,d});
	lim = (n / 4);
	poly g;g.resize(lim + 1);
	for(int i = 0;i <= lim;i++)g[i] = ::g(i);
	poly f;f.resize(lim + 1);
	for(int x = 0;x <= lim;x++)
		for(int i = x;i <= lim;i++)
			f[x] = add(f[x],mul(mul(sgn(i - x),binom[i][x]),g[i]));
	int ans = g[0];
	for(int i = 1;i <= lim;i++)ans = sub(ans,f[i]);
	printf("%d
",ans);
	return 0;
}