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

阿新 • • 發佈：2020-12-23

點此看題面

有一個長度為$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;
}