【组合计数+NTT优化卷积】BZOJ5306 [HAOI2018] 染色

【题目】
lydsy
一个长度为 $n$ 的序列，每个位置可以被染成 $m$ 种颜色中的一种。若一种方案中出现次数恰好为 $S$ 的颜色数有 $K$ 种，则会有 $W_K$ 的愉悦值。问所有方案的愉悦值总和对 $1004535809$ 取模的结果。

$n\leq 10^7,m\leq 10^5,S\leq 150$

【解题思路】
首先这个模数就很 $\text{NTT}$ 了（ $1004535809=2^{21}\cdot 479+1$ ，原根为 $3$ ）

然后我就发现我不会做了。于是以下来自各种 $\text{blog}$

设 $N=\min(m,\lfloor \frac n S \rfloor)$ 表示出现 $S$ 次的颜色数上界，设 $f_i$ 表示染色后出现了 $S$ 次的颜色有 $i$ 种的方案数，则答案就是：
$\sum_{i=0}^N w_i\cdot f_i$
这个 $f_i$ 直接求比较难求，于是考虑这样一个容斥，设 $g_i=i!$ ：
$f_i={m\choose i}{n \choose iS} \frac {(g_{iS})!} {g_S^i} \sum_{j=0}^{N-i} (-1)^j {m-i\choose j}{n-iS\choose jS} \frac {g_{jS}} {g_S^j} (m-i-j)^{n-iS-jS}$
这个柿子的意思大概是先从 $m$ 种颜色中选出 $i$ 种。然后从这 $n$ 个位置中选出 $iS$ 个进行可重排列，就是 ${n \choose iS} \frac {(g_{iS})!} {g_S^i}$ 了。接下来剩下的 $m-i$ 种颜色，在 $n-iS$ 个位置上随便填，但由于可能会使某种颜色出现了 $S$ 次，所以后面这个东西我们需要容斥。容斥的意义就是额外出现了 $S$ 次的颜色有 $j$ 种，我们就需要在 $m-i$ 种颜色中选出这 $j-i$ 种，并在 $n-iS$ 个位置中选出 $jS$ 个进行可重排列，剩下的位置随便填。

这样 $(m-i-j)^{n-iS-jS}$ 不太好处理，不妨将第二个和式中的 $j$ 意义改变为实际上有 $j$ 种颜色出现了 $S$ 次，那么就是用 $j-i$ 替换 $j$ 并改变循环下标：
$f_i={m\choose i}{n \choose iS} \frac {(g_{iS})!} {g_S^i} \sum_{j=i}^{N} (-1)^{j-i} {m-i\choose j-i}{n-iS\choose jS-iS} \frac {g_{jS-iS}} {g_S^{j-i}} (m-j)^{n-jS}$
接下来就是化简这个柿子，那么将组合数写成阶乘形式后化简：
$f_i=\frac {g_m\cdot g_n} {g_i}\sum_{j=i}^N\frac {(-1)^{j-i}\cdot (m-j)^{n-jS}}{g_{j-i}\cdot g_{m-j}\cdot g_{n-jS}\cdot g_S^j}$
这里面的 $j$ 不太好做，不妨把 $j$ 提到外层，那么最后的答案可以得到：
$\begin{aligned} ans &=\sum_{i=0}^N w_i\cdot f_i \\ &=\sum_{i=0}^Nw_i\cdot \frac {g_m\cdot g_n} {g_i}\sum_{j=i}^N\frac {(-1)^{j-i}\cdot (m-j)^{n-jS}}{g_{j-i}\cdot g_{m-j}\cdot g_{n-jS}\cdot g_S^j} \\ &=g_m\cdot g_n\cdot \sum_{j=0}^N\frac {(m-j)^{n-jS}}{g_{m-j}\cdot g_{n-jS}\cdot g_S^j} \sum_{i=0}^j \frac {(-1)^{j-i}\cdot w_i}{g_i\cdot g_{j-i}} \end{aligned}$
现在左边是一个只与 $j$ 有关的柿子，右边是 $A(x)=\frac {w_x} {g_x}$ 和 $B(x)=\frac {(-1)^x}{g_x}$ 的卷积。

那么我们就可以愉快地做 $\text{NTT}$ 了，最后复杂度 $O(n+m\log m)$

【参考代码】~~荣获BZOJ垫底~~

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

typedef long long ll;
const int N=1e7+10,M=262333,mod=1004535809,G=3;

namespace IO
{
	int read()
	{
		int ret=0;char c=getchar();
		while(!isdigit(c)) c=getchar();
		while(isdigit(c)) ret=ret*10+(c^48),c=getchar();
		return ret;
	}
	void write(int x){if(x>9)write(x/10);putchar(x%10^48);}
	void writeln(int x){write(x);putchar('\n');}
}
using namespace IO;

namespace Math
{
	int fac[N],ifac[N];
	int upm(int x){return x>=mod?x-mod:(x<0?x+mod:x);}
	void up(int &x,int y){x=upm(x+y);}
	int mul(int x,int y){return 1ll*x*y%mod;}
	int qpow(int x,int y)
	{
		int res=1;
		for(;y;y>>=1,x=mul(x,x)) if(y&1) res=mul(res,x);
		return res;
	}
	int C(int x,int y){if(y>x)return 0;return mul(fac[x],mul(ifac[x-y],ifac[y]));}
	void initsth()
	{
		fac[0]=1;for(int i=1;i<N;++i)fac[i]=mul(fac[i-1],i);
		ifac[N-1]=qpow(fac[N-1],mod-2);for(int i=N-2;~i;--i)ifac[i]=mul(ifac[i+1],i+1);
	}
}
using namespace Math;

namespace Poly
{
	int m,L;
	int rev[M];
	void ntt(int *a,int n,int op)
	{
		for(int i=0;i<n;++i) if(i<rev[i]) swap(a[i],a[rev[i]]);
		for(int i=1;i<n;i<<=1)
		{
			int wn=qpow(G,(mod-1)/(i<<1));
			if(!~op) wn=qpow(wn,mod-2);
			for(int j=0;j<n;j+=(i<<1))
			{
				int w=1;
				for(int k=0;k<i;++k,w=mul(w,wn))
				{
					int x=a[j+k],y=mul(w,a[i+j+k]);
					a[j+k]=upm(x+y);a[i+j+k]=upm(x-y);
				}
			}
		}
		if(!~op) for(int inv=qpow(n,mod-2),i=0;i<n;++i) a[i]=mul(a[i],inv);
	}
	void mulpoly(int *a,int *b,int *c,int n)
	{
		for(m=1,L=0;m<=n;m<<=1,++L);
		for(int i=0;i<m;++i) rev[i]=(rev[i>>1]>>1)|((i&1)<<(L-1));
		ntt(a,m,1);ntt(b,m,1);
		for(int i=0;i<m;++i) c[i]=mul(a[i],b[i]);
		ntt(c,m,-1);
	}
}
using namespace Poly;

namespace DreamLolita
{
	int n,m,S,lim;
	int w[M],f[M],g[M];
	void input()
	{
		n=read();m=read();S=read();lim=min(m,n/S);
		for(int i=0;i<=m;++i) w[i]=read();
	}
	void initf()
	{	
	 	for(int i=0;i<=lim;++i) f[i]=mul(w[i],ifac[i]),g[i]=upm((i&1?-1:1)*ifac[i]);
	 	mulpoly(f,g,f,lim*2);
	}
	void getans()
	{
		int ans=0;
		for(int j=0;j<=lim;++j)
		{
			int t1=qpow(m-j,n-j*S),t2=mul(ifac[m-j],mul(ifac[n-j*S],qpow(ifac[S],j)));
			up(ans,mul(mul(t1,t2),f[j]));
		}
		ans=mul(ans,mul(fac[m],fac[n]));
		writeln(ans);
	}
	void solution(){initsth();input();initf();getans();}
}

int main()
{
#ifdef Durant_Lee
	freopen("BZOJ5306.in","r",stdin);
	freopen("BZOJ5306.out","w",stdout);
#endif
	DreamLolita::solution();
	return 0;
}

【组合计数+NTT优化卷积】BZOJ5306 [HAOI2018] 染色

猜你喜欢