【JZOJ 5559】花朵

Description

Solution

暴力$O(n^3)$显然，还可以用NTT优化，
链的做法也简单，先设左/右有没有选，再设选了几个，用分治NTT优化成$n\log(n)^2$
菊花更简单，直接分治NTT，

那树怎么做呢？
链剖，长链上的用链的做法，链顶再用菊花的做法合并，

复杂度：$O(n\log(n)^3)$

Code

#include <cstdio>
#include <algorithm>
#define fo(i,a,b) for(int i=a;i<=b;++i)
#define fod(i,a,b) for(int i=a;i>=b;--i)
#define efo(i,q) for(int i=A[q];i;i=B[i][0])
#define min(q,w) ((q)>(w)?(w):(q))
#define max(q,w) ((q)<(w)?(w):(q))
using namespace std;
typedef long long LL;
const int N=100500,mo=998244353;
int read(int &n)
{
    char ch=' ';int q=0,w=1;
    for(;(ch!='-')&&((ch<'0')||(ch>'9'));ch=getchar());
    if(ch=='-')w=-1,ch=getchar();
    for(;ch>='0' && ch<='9';ch=getchar())q=q*10+ch-48;n=q*w;return n;
}
int m,n,ans;
int B[2*N][2],A[N],B0;
void link(int q,int w)
{
    B[++B0][0]=A[q],A[q]=B0,B[B0][1]=w;
    B[++B0][0]=A[w],A[w]=B0,B[B0][1]=q;
}
LL ksm(LL q,int w)
{
    LL ans=1;
    for(;w;w>>=1,q=q*q%mo)if(w&1)ans=ans*q%mo;
    return ans;
}
LL c[N*3],c1[N*3],c2[N*3];
int W[1048576+10],W0=1048576;
void DFT(LL *a,int n,int ws,int K)
{
    fo(i,0,n-1)
    {
        int q=0;
        for(int j=ws,w=i;j;w>>=1,--j)q=(q<<1)+(w&1);
        c[q]=a[i];
    }
    for(int I=2;I<=n;I<<=1)
    {
        register int mid=I>>1;
        for(int i=0;i<n;i+=I)
        {
            fo(j,0,mid-1)
            {
                LL w=W[K>0?W0/I*j:W0-W0/I*j];
                LL t=w*c[i+j+mid]%mo;
                c[i+j+mid]=(c[i+j]-t)%mo;
                c[i+j]=(c[i+j]+t)%mo;
            }
        }
    }
    if(K<0)
    {
        LL t=ksm(n,mo-2);
        fo(i,0,n-1)c[i]=c[i]*t%mo;
    }
}
void NTT(LL *a,LL *b,int n,int K=0)
{
    int m=1,ws=1;
    for(;m<=n;m<<=1,++ws);
    m<<=1;
    if(!K)fo(i,0,n)c1[i]=a[i],c2[i]=b[i];
    fo(i,n+1,m)c1[i]=c2[i]=0;
    DFT(c1,m,ws,1);
    fo(i,0,m-1)c1[i]=c[i];
    DFT(c2,m,ws,1);
    fo(i,0,m-1)c1[i]=c1[i]*c[i]%mo;
    DFT(c1,m,ws,-1);
}

struct qqww
{
    int v,si,lt,fa,zx,nx;
}a[N];

int d[N],d1[N];
LL F[20][4][N];
int f[N],g[N],fv[N];//YES  NO
int BF[100*N][2],BF0;
void divide(int l,int r,int C,int Alln)
{
    if(l>r)return;
    if(l==r)
    {
        for(int q=f[d[l]],i=0;q;++i,q=BF[q][0])F[C][3][i]=BF[q][1];
        for(int q=g[d[l]],i=0;q;++i,q=BF[q][0])F[C][0][i]=BF[q][1];
        return;
    }
    int mid=l,mid1=-100,mid2=0,t=0;
    fo(i,l,r)
    {
        t+=d1[i];
        if(min(t,Alln-t)>mid1)mid1=min(t,Alln-t),mid2=t,mid=i;
    }
    divide(l,mid,C,mid2);
    divide(mid+1,r,C+1,Alln-mid2);
    mid=min(max(Alln-mid2,mid2),m);
    mid1=min(Alln,m);
    fo(L,0,1)fo(R,0,1)
    {
        fo(i,0,mid)F[0][L*2+R][i]=(F[C+1][2+R][i]+F[C+1][R][i])%mo;
        NTT(F[C][L*2],F[0][L*2+R],mid);
        fo(i,0,mid1)F[0][L*2+R][i]=c[i];
        NTT(F[C][L*2+1],F[C+1][R],mid);
        fo(i,0,mid1)F[0][L*2+R][i]=(F[0][L*2+R][i]+c[i])%mo;
    }
    fo(j,0,3)fo(i,0,mid1)F[C][j][i]=F[0][j][i],F[C+1][j][i]=F[0][j][i]=0;
}
void divide1(int l,int r,int C,int Alln)
{
    if(l>r)return;
    if(l==r)
    {
        for(int q=f[d[l]],i=0;q;++i,q=BF[q][0])F[C][0][i]=BF[q][1];
        for(int q=g[d[l]],i=0;q;++i,q=BF[q][0])F[C][1][i]=BF[q][1];
        return;
    }
    int mid=l,mid1=-100,mid2=0,t=0;
    fo(i,l,r)
    {
        t+=d1[i];
        if(min(t,Alln-t)>mid1)mid1=min(t,Alln-t),mid2=t,mid=i;
    }
    divide1(l,mid,C,mid2);
    divide1(mid+1,r,C+1,Alln-mid2);
    mid=min(max(Alln-mid2,mid2),m);
    mid1=min(Alln,m);
    fo(K,0,1)
    {
        NTT(F[C][K],F[C+1][K],mid);
        fo(i,0,mid1)F[C][K][i]=c[i],F[C+1][K][i]=0;
    }
}


int dfsSi(int q,int fa)
{
    efo(i,q)if(B[i][1]!=fa)a[q].si+=dfsSi(B[i][1],q);
    return ++a[q].si;
}
void dfsf(int q,int fa,int lt)
{
    a[q].fa=fa;
    a[q].lt=(lt?lt:(lt=q));
    int mx=0;
    efo(i,q)if(B[i][1]!=fa&&a[B[i][1]].si>a[mx].si)mx=B[i][1];
    a[q].nx=mx;
    if(mx)dfsf(mx,q,lt);
    efo(i,q)if(B[i][1]!=fa&&B[i][1]!=mx)dfsf(B[i][1],q,0);
}
void dfs(int q,int fa)
{
    BF[g[q]=++BF0][1]=1;fv[q]=1;
    BF[++BF0][1]=a[q].v;BF[BF0+1][0]=BF0;BF[f[q]=++BF0][1]=0;
    efo(i,q)if(B[i][1]!=fa)dfs(B[i][1],q);
    int t=1;
    d[0]=1;d[1]=q;d1[1]=1;
    efo(i,q)if(B[i][1]!=fa&&B[i][1]!=a[q].nx)
        d[++d[0]]=B[i][1],t+=d1[d[0]]=fv[B[i][1]];
    divide1(1,d[0],0,t);
    t=min(m,t);g[q]=f[q]=0;fv[q]=t;
    fod(i,t,0)BF[++BF0][0]=g[q],g[q]=BF0,BF[BF0][1]=F[0][1][i];
    fod(i,t,0)BF[++BF0][0]=f[q],f[q]=BF0,BF[BF0][1]=F[0][0][i];
    fo(j,0,1)fo(i,0,t)F[0][j][i]=0;
    if(a[q].lt==q)
    {
        // printf("%d\n",q);
        int w=q,t=0;
        for(d[0]=0;w;d[++d[0]]=w,t+=d1[d[0]]=fv[w],w=a[w].nx);
        divide(1,d[0],1,t);
        int T=min(m,t);fv[q]=T;
        g[q]=f[q]=0;
        fod(i,T,0)BF[++BF0][0]=g[q],g[q]=BF0,BF[BF0][1]=(F[1][2][i]+F[1][3][i]+F[1][0][i]+F[1][1][i])%mo;
        fod(i,T,0)BF[++BF0][0]=f[q],f[q]=BF0,BF[BF0][1]=(F[1][0][i]+F[1][1][i])%mo;
        fo(j,0,3)fo(i,0,T)F[1][j][i]=0;
    }
}
int main()
{
    freopen("flowers.in","r",stdin);
    freopen("flowers.out","w",stdout);
    int q,w;
    W[0]=1;W[1]=ksm(3,(mo-1)/W0);
    fo(i,2,W0)W[i]=(LL)W[i-1]*W[1]%mo;
    read(n),read(m);
    fo(i,1,n)read(a[i].v);
    fo(i,1,n-1)read(q),read(w),link(q,w);
    dfsSi(1,0);
    dfsf(1,0,0);
    dfs(1,0);
    for(q=g[1];BF[q][0];q=BF[q][0]);
    printf("%d\n",(BF[q][1]+mo)%mo);
    return 0;
}