【洛谷4491】[HAOI2018] 染色（二项式反演+NTT）

点此看题面

• 有一个长度为\(n\)的序列和\(m\)种颜色，给定\(S\)以及一个序列\(W\)。
• 对于一种染色方案，假设其中有\(k\)种颜色恰好出现了\(S\)次，则其价值为\(W_k\)。
• 求所有染色方案的价值之和。
• \(n\le10^7,m\le10^5,S\le150\)

二项式反演

考虑设\(f_i\)表示恰好出现\(S\)次的颜色至少有\(i\)种的方案数。

首先，从\(m\)种颜色中选\(i\)种颜色的方案数为\(C_m^i\)。

其次，从\(n\)个位置中我们要选出\(i\times S\)个位置染上这\(i\)种颜色，剩余位置染上剩余\(m-i\)种颜色，这里的选择方案数位\(C_n^{i\times S}\)。

对于这\(i\)种颜色，染色方案是一个可重全排列，即\(\frac{(i\times S)!}{(S!)^i}\)。

对于剩余\(m-i\)种颜色，可以随意染色，染色方案就是\((m-i)^{n-i\times S}\)。

综上，得出计算式：

\[f_i=C_m^i\times C_n^{i\times S}\times \frac{(i\times S)!}{(S!)^i}\times (m-i)^{n-i\times S} \]

然后我们设\(g(i)\)表示恰好出现\(S\)次的颜色恰好有\(i\)种的方案数。

令\(l=\min\{m,\lfloor\frac nS\rfloor\}\)，利用组合数，有一个显然的关系式：

\[f_i=\sum_{j=i}^lC_j^i\times g_j \]

对于这个式子考虑二项式反演，得到：

\[g_i=\sum_{j=i}^l(-1)^{j-i}\times C_j^i\times f_j \]

\(NTT\)优化

考虑转化一下这个式子，首先拆开组合数并移项：

\[g_i\times i!=\sum_{j=i}^l\frac{(-1)^{j-i}}{(j-i)!}\times (f_j\times j!) \]

构造两个多项式\(A(x)=\sum_{i=0}^{l}(f_i\times i!)x^i,B(x)=\sum_{i=0}^l\frac{(-1)^{l-i}}{(l-i)!}x^i\)（注意\(B(x)\)的系数经过了翻转，以形成一个卷积的形式）

令\(C(x)=A(x)*B(x)\)，则\(C(x)\)的\(n+i\)次项系数就是\(g_i\times i!\)。

最后答案的统计应该是非常简单的。

代码：\(O(nlogn)\)

#include<bits/stdc++.h>
#define Tp template<typename Ty>
#define Ts template<typename Ty,typename... Ar>
#define Reg register
#define RI Reg int
#define Con const
#define CI Con int&
#define I inline
#define W while
#define N 10000000
#define M 100000
#define X 1004535809
#define C(x,y) (1LL*Fac[x]*IFac[y]%X*IFac[(x)-(y)]%X)
using namespace std;
int n,m,S,a[M+5],f[M+5],Fac[N+5],IFac[N+5],l,A[N<<2],B[N<<2];
I int QP(RI x,RI y) {RI t=1;W(y) y&1&&(t=1LL*t*x%X),x=1LL*x*x%X,y>>=1;return t;}
namespace Poly//多项式
{
	#define PR 3
	int P,L,R[N<<2];I void NTT(int* s,CI op)//NTT
	{
		RI i,j,k,x,y,U,S;for(i=0;i^P;++i) i<R[i]&&(x=s[i],s[i]=s[R[i]],s[R[i]]=x);
		for(i=1;i^P;i<<=1) for(U=QP(QP(PR,op),(X-1)/(i<<1)),j=0;j^P;j+=i<<1) for(S=1,k=0;
			k^i;++k,S=1LL*S*U%X) s[j+k]=((x=s[j+k])+(y=1LL*S*s[i+j+k]%X))%X,s[i+j+k]=(x-y+X)%X;
	}
	I void Mul(CI n,int* A,int* B)//卷积
	{
		RI i;P=1,L=0;W(P<=(n<<1)) P<<=1,++L;for(i=0;i^P;++i) R[i]=((R[i>>1]>>1)|((i&1)<<L-1));
		for(NTT(A,1),NTT(B,1),i=0;i^P;++i) A[i]=1LL*A[i]*B[i]%X;
		RI t=QP(P,X-2);for(NTT(A,X-2),i=0;i<=2*n;++i) A[i]=1LL*A[i]*t%X;
	}
}
int main()
{
	RI i;for(scanf("%d%d%d",&n,&m,&S),l=min(m,n/S),i=0;i<=n;++i) scanf("%d",a+i);
	RI Mx=max(n,m);for(Fac[0]=i=1;i<=Mx;++i) Fac[i]=1LL*Fac[i-1]*i%X;//预处理阶乘
	for(IFac[Mx]=QP(Fac[Mx],X-2),i=Mx;i;--i) IFac[i-1]=1LL*IFac[i]*i%X;//预处理阶乘逆元
	for(i=0;i<=l;++i) f[i]=1LL*C(m,i)*C(n,i*S)%X*Fac[i*S]%X*QP(IFac[S],i)%X*QP(m-i,n-i*S)%X;//计算f[i]
	for(i=0;i<=l;++i) A[i]=1LL*f[i]*Fac[i]%X,B[l-i]=1LL*(i&1?X-1:1)*IFac[i]%X;//构造两个多项式
	RI t=0;for(Poly::Mul(l,A,B),i=0;i<=l;++i) t=(1LL*A[l+i]*IFac[i]%X*a[i]+t)%X;//卷积，然后统计答案
	return printf("%d\n",t),0;
}