【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;
}