【XSY2666】排列问题 DP 容斥原理 分治FFT
题目大意
有\(n\)种颜色的球，第\(i\)种有\(a_i\)个。设\(m=\sum a_i\)。你要把这\(m\)个小球排成一排。有\(q\)个询问，每次给你一个\(x\)，问你有多少种方案使得相邻的小球同色的对数为\(x\)。
\(n\leq 10000,m\leq 200000\)
题解
我们考虑把这些小球分段，每段内所有小球颜色相同，但相邻两段的小球颜色可以相同。
设第\(i\)种颜色有\(b_i\)段，那么分\(j\)段的方案数是\(\frac{(\sum b_i)!}{\sum(bi!)}=\frac{j!}{\sum(bi!)}\)
那么先DP，设\(f_{i,j}\)为前\(i\)种颜色，分了\(j\)段的方案数\(\div b_i!\)。显然枚举第\(i\)中颜色分\(k\)段得
\[f_{i,j}+=f_{i-1,j-k}\times \binom{a_i-1}{k-1}\times\frac{1}{k!}
\]
那个组合数是插板法得到的。
这个DP的时间复杂度是\(O(m^2)\)（因为枚举第\(i\)种颜色时\(k=1\ldots a_i,j=1\ldots s_i\)（\(s\)为\(a\)的前缀和））
然后这个东西可以分治FFT优化到\(O(m\log m\log n)\)
这样我们得到了分成\(i\)段的方案数\(g_i=f_{n,i}\times i!\)，但相邻两段可能颜色相同。我们还要减掉这种情况。
就是对于一种实际上分成 \(j\) 段的方案，它在分成 \(i\) 段的方案数中会被计算 \(\binom{m-j}{m-i}\) 次（就是在 \(m-j\) 个间隔中取 \(m-i\) 个）。
答案 \(ans_i=g_i-\sum_{j<i}ans_j\binom{m-j}{i-j}\)
可以简单暴力的通过分治FFT优化到\(O(m\log^2 m)\)。但有更好的做法。
考虑容斥。其实总的\(g_j\)对\(ans_i\)的贡献就是\({(-1)}^{i-j}\binom{m-j}{i-j}\)。直接FFT一次就可以得到答案。
\[\begin{align}
ans_{k->i}&=\sum_{j=k}^{i-1}{(-1)^{j-k}}\binom{m-k}{j-k}\binom{m-j}{i-j}\\
&=\sum_{j=k}^{i-1}{(-1)^{j-k}}\frac{(m-k)!(m-j)!}{(j-k)!(m-j)!(i-j)!(m-i)!}\\
&=\sum_{j=k}^{i-1}{(-1)^{j-k}}\frac{(m-k)!}{(j-k)!(i-j)!(m-i)!}\\
&=\frac{(m-k)!}{(m-i)!(i-k)!}\sum_{j=k}^{i-1}{(-1)^{j-k}}\frac{(i-k)!}{(i-j)!(j-k)!}\\
&=\binom{m-k}{i-k}\sum_{j=k}^{i-1}{(-1)^{j-k}}\binom{i-k}{j-k}\\
&=\binom{m-k}{i-k}{(-1)}^{i-k}
\end{align}
\]
那么相邻的小球同色的对数为\(x\)的答案就是\(ans_{m-x}\)。
时间复杂度：\(O(m\log m\log n+q)\)
代码
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<cstdlib>
#include<ctime>
#include<utility>
#include<cmath>
#include<functional>
#include<vector>
#include<queue>
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
typedef pair<int,int> pii;
typedef pair<ll,ll> pll;
void sort(int &a,int &b)
{
if(a>b)
swap(a,b);
}
void open(const char *s)
{
#ifndef ONLINE_JUDGE
char str[100];
sprintf(str,"%s.in",s);
freopen(str,"r",stdin);
sprintf(str,"%s.out",s);
freopen(str,"w",stdout);
#endif
}
int rd()
{
int s=0,c;
while((c=getchar())<'0'||c>'9');
do
{
s=s*10+c-'0';
}
while((c=getchar())>='0'&&c<='9');
return s;
}
void put(int x)
{
if(!x)
{
putchar('0');
return;
}
static int c[20];
int t=0;
while(x)
{
c[++t]=x%10;
x/=10;
}
while(t)
putchar(c[t--]+'0');
}
int upmin(int &a,int b)
{
if(b<a)
{
a=b;
return 1;
}
return 0;
}
int upmax(int &a,int b)
{
if(b>a)
{
a=b;
return 1;
}
return 0;
}
const int p=998244353;
int fp(int a,int b)
{
int s=1;
for(;b;b>>=1,a=1ll*a*a%p)
if(b&1)
s=1ll*s*a%p;
return s;
}
int inv[600010];
int fac[600010];
int ifac[600010];
namespace ntt
{
const int g=3;
int rev[600010];
int w1[600010];
int w2[600010];
int n;
void init(int m)
{
n=1;
while(n<=m)
n<<=1;
int i;
rev[0]=0;
for(i=1;i<n;i++)
rev[i]=(rev[i>>1]>>1)|(i&1?n>>1:0);
for(i=1;i<=n;i<<=1)
{
w1[i]=fp(g,(p-1)/i);
w2[i]=fp(w1[i],p-2);
}
}
void ntt(int *a,int t)
{
int i,j,k;
int u,v,w,wn;
for(i=0;i<n;i++)
if(rev[i]<i)
swap(a[i],a[rev[i]]);
for(i=2;i<=n;i<<=1)
{
wn=(t==1?w1[i]:w2[i]);
for(j=0;j<n;j+=i)
{
w=1;
for(k=j;k<j+i/2;k++)
{
u=a[k];
v=1ll*a[k+i/2]*w%p;
a[k]=(u+v)%p;
a[k+i/2]=(u-v)%p;
w=1ll*w*wn%p;
}
}
}
if(t==-1)
{
int inv=fp(n,p-2);
for(i=0;i<n;i++)
a[i]=1ll*a[i]*inv%p;
}
}
};
int g[600010];
int h[600010];
int ans[600010];
int a[600010];
int s[600010];
int n,m;
void add(int &a,int b)
{
a=(a+b)%p;
}
typedef vector<int> vec;
vec mul(vec &a,vec &b)
{
static int c[600010],d[600010];
int n1=a.size()-1;
int n2=b.size()-1;
int m=n1+n2+1;
ntt::init(m);
int i;
for(i=0;i<=n1;i++)
c[i]=a[i];
for(i=n1+1;i<ntt::n;i++)
c[i]=0;
for(i=0;i<=n2;i++)
d[i]=b[i];
for(i=n2+1;i<ntt::n;i++)
d[i]=0;
ntt::ntt(c,1);
ntt::ntt(d,1);
for(i=0;i<ntt::n;i++)
c[i]=1ll*c[i]*d[i]%p;
ntt::ntt(c,-1);
vec s(n1+n2+1);
for(i=1;i<=n1+n2;i++)
s[i]=c[i];
return s;
}
vec solve(int l,int r)
{
if(l==r)
{
vec s(a[l]+1);
int i;
for(i=1;i<=a[l];i++)
s[i]=1ll*ifac[i-1]*ifac[i]%p*ifac[a[l]-i]%p;
return s;
}
int mid=(l+r)>>1;
vec s1=solve(l,mid);
vec s2=solve(mid+1,r);
return mul(s1,s2);
}
int c[600010];
int d[600010];
priority_queue<pii,vector<pii>,greater<pii> > q;
void gao()
{
int i;
c[0]=0;
for(i=1;i<=m;i++)
c[i]=g[i];
for(i=0;i<=m;i++)
{
d[i]=ifac[i];
if(i&1)
d[i]=-d[i];
}
ntt::init(2*m);
for(i=m+1;i<ntt::n;i++)
c[i]=d[i]=0;
ntt::ntt(c,1);
ntt::ntt(d,1);
for(i=0;i<ntt::n;i++)
c[i]=1ll*c[i]*d[i]%p;
ntt::ntt(c,-1);
for(i=1;i<=m;i++)
g[i]=c[i];
}
int t=0;
vec f[20010];
int main()
{
open("c");
scanf("%d",&n);
int i;
for(i=1;i<=n;i++)
{
scanf("%d",&a[i]);
s[i]=s[i-1]+a[i];
}
m=s[n];
inv[0]=inv[1]=fac[0]=fac[1]=ifac[0]=ifac[1]=1;
for(i=2;i<=m;i++)
{
inv[i]=-1ll*p/i*inv[p%i]%p;
#ifndef ONLINE_JUDGE
inv[i]=(inv[i]+p)%p;
#endif
fac[i]=1ll*fac[i-1]*i%p;
ifac[i]=1ll*ifac[i-1]*inv[i]%p;
}
// f[0][0]=1;
int times=1;
for(i=1;i<=n;i++)
times=1ll*times*fac[a[i]-1]%p;
// for(i=1;i<=n;i++)
// {
// times=times*fac[a[i]-1]%p;
// for(j=1;j<=s[i];j++)
// {
// for(k=1;k<=a[i]&&k<=j;k++)
// add(f[i][j],f[i-1][j-k]*ifac[k-1]%p*ifac[a[i]-k]%p*ifac[k]%p);
//// add(f[i][j],f[i-1][j-k]*c(a[i]-1,k-1)%p*ifac[k]%p);
//// f[i][j]=f[i][j]*fac[a[i]-1]%p;
// }
// }
int j;
for(i=1;i<=n;i++)
{
f[i].resize(a[i]+1);
for(j=1;j<=a[i];j++)
f[i][j]=1ll*ifac[j-1]*ifac[j]%p*ifac[a[i]-j]%p;
q.push(pii(a[i],i));
}
t=n;
for(i=1;i<n;i++)
{
int n1=q.top().first;
int x=q.top().second;
q.pop();
int n2=q.top().first;
int y=q.top().second;
q.pop();
f[++t]=mul(f[x],f[y]);
f[x].clear();
f[y].clear();
q.push(pii(n1+n2+1,t));
}
vec ss=f[t];
// vec ss=solve(1,n);
for(i=1;i<=m;i++)
g[i]=1ll*ss[i]*fac[i]%p*times%p;
#ifndef ONLINE_JUDGE
for(i=1;i<=m;i++)
add(g[i],p);
#endif
// g[i]=f[n][i]*fac[i]%p*times%p;
for(i=1;i<=m;i++)
g[i]=1ll*g[i]*fac[m-i]%p;
gao();
for(i=1;i<=m;i++)
{
g[i]=1ll*g[i]*ifac[m-i]%p;
add(g[i],p);
}
// for(i=1;i<=m;i++)
// {
// for(j=1;j<i;j++)
// add(ans[i],h[j]%p*ifac[i-j]%p);
// ans[i]=-ans[i]*ifac[m-i]%p;
// ans[i]=(ans[i]+g[i])%p;
// add(ans[i],-ans[j]*c(m-j,i-j));
// add(ans[i],p);
// h[i]=ans[i]*fac[m-i]%p;
// }
int q;
int x;
scanf("%d",&q);
while(q--)
{
scanf("%d",&x);
printf("%lld\n",g[m-x]);
}
return 0;
}
