题解:AT_pakencamp_2024_day2_b Biscuit and Ants

· · 题解

啊,饼干。无疑所有人都喜欢饼干;但是,并不是所有人都喜欢蚂蚁。其实,你知道吗,蚂蚁喜欢饼干。

从拆贡献开始。只考虑点 u 的子树,它的贡献是 f_u,也就是能够使该点最后有饼干的初始方案数。

  1. u 初始有饼干:f_u=2^{z_u-1},其中 z_u 表示子树大小。

你需要以每个点为根都进行全树 DP 才可以。我想我们得到了一个 \Theta(n^2m) 的做法。不过,2 类转移是可撤销的,所以换根 DP 绝非困难。现在做到 \Theta(nm) 了。

不过,你大概已经注意到了。如果 d_v<m,这个点的初始状态就是最终状态,那么 f_v=2^{z_v}-f_v=2^{z_v-1}。这类子树转移很特殊,将其分离。设 X=\sum\limits_{v\in\operatorname{son}(u),d_v<m}(z_v-1)Y=\sum\limits_{v\in\operatorname{son}(u)}[d_v<m],在仅对 d_v\ge m 的子节点完成暴力转移后,事实上:f_u=2^{z_u-1}+2^X\sum\limits_{i=\max(0,m-Y)}^ms_{u,i}\sum\limits_{j=m-i}^Y\binom{Y}{j},即额外取 j 个满足 d_v<m 的点上的饼干。如果预处理了组合数的前缀和,f_u 可以被 \Theta(m) 地求出。

我想,换根撤销并不难。

由于 \sum d=2(n-1),满足 d_u\ge m 的点数只有 \lfloor\frac{n}{m}\rfloor,每个点需要 \Theta(m) 复杂度,总共是 \Theta(n)

#include<bits/stdc++.h>
#define LL long long
#define fi first
#define se second
using namespace std;

const int N=1e6+10;
int tid,n,m,mod,f[N],z[N],pw[N],ans,g[N],fct[N],inv[N],d[N],X[N],Y[N];
vector<int>e[N],s[N],c[N];

int rd(){
    int x=0;
    char ch=getchar();
    while(!isdigit(ch))ch=getchar();
    while(isdigit(ch))x=x*10+ch-'0',ch=getchar();
    return x;
}
int qpow(int x,int y){
    int ret=1;
    while(y){
        if(y&1)ret=1ll*ret*x%mod;
        y>>=1;x=1ll*x*x%mod;
    }
    return ret;
}
void add(int &x,int y){if((x+=y)>=mod)x-=mod;}
int C(int x,int y){
    if(x<y||y<0)return 0;
    return 1llu*fct[x]*inv[x-y]%mod*inv[y]%mod;
}
void dfs1(int u,int fa){
    z[u]=1;
    for(int v:e[u])if(v!=fa)dfs1(v,u),z[u]+=z[v];
    f[u]=pw[z[u]-1];
    if(e[u].size()<m)return;
    s[u].resize(m+1,0);s[u][0]=1;
    for(int v:e[u])if(v!=fa)
        if(e[v].size()>=m)for(int i=m;i>=0;i--){
            s[u][i]=1llu*(i<m?pw[z[v]]-f[v]+mod:pw[z[v]])*s[u][i]%mod;
            if(i)add(s[u][i],1llu*s[u][i-1]*f[v]%mod);
        }else Y[u]++,X[u]+=z[v]-1;
    c[u].resize(Y[u]+1,0);c[u][0]=1;
    for(int i=1;i<=Y[u];i++)c[u][i]=(C(Y[u],i)+c[u][i-1])%mod;
    for(int i=max(0,m-Y[u]);i<=m;i++)
        add(f[u],1llu*(c[u][Y[u]]-(i<m?c[u][m-i-1]:0)+mod)*pw[X[u]]%mod*s[u][i]%mod);
}
void dfs2(int u,int fa){
    if(e[u].size()<m){
        for(int v:e[u])if(v!=fa)g[u]=pw[n-1-z[v]],dfs2(v,u);
        g[u]=pw[n-1];
        return;
    }
    vector<int>t(m+1,0);
    if(u!=1)for(int i=m;i>=0;i--){
        s[u][i]=1llu*(i<m?pw[n-z[u]]-g[fa]+mod:pw[n-z[u]])*s[u][i]%mod;
        if(i)add(s[u][i],1llu*s[u][i-1]*g[fa]%mod);
    }
    for(int v:e[u])if(v!=fa){
        if(e[v].size()>=m){
            int j1=qpow(pw[z[v]]-f[v]+mod,mod-2),j2=qpow(pw[z[v]],mod-2);
            for(int i=0;i<=m;i++){
                t[i]=s[u][i];
                if(i)add(t[i],mod-1llu*t[i-1]*f[v]%mod);
                t[i]=1llu*(i<m?j1:j2)*t[i]%mod;
            }
            g[u]=pw[n-1-z[v]];
            for(int i=max(0,m-Y[u]);i<=m;i++)
                add(g[u],1llu*(c[u][Y[u]]-(i<m?c[u][m-i-1]:0)+mod)*pw[X[u]]%mod*t[i]%mod);
        }
        dfs2(v,u);
    }
    g[u]=pw[n-1];
    for(int i=max(0,m-Y[u]);i<=m;i++)
        add(g[u],1llu*(c[u][Y[u]]-(i<m?c[u][m-i-1]:0)+mod)*pw[X[u]]%mod*s[u][i]%mod);
}

signed main(){
    scanf("%d%d%d",&n,&m,&mod);
    if(m==1){
        ans=1;
        for(int i=1;i<=n;i++)ans=2*ans%mod;
        printf("%d\n",1llu*(ans-1)*n%mod);
        return 0;
    }
    for(int i=1,x,y;i<n;i++)x=rd(),y=rd(),e[x].push_back(y),e[y].push_back(x);
    pw[0]=fct[0]=1;
    for(int i=1;i<=n;i++)pw[i]=2*pw[i-1]%mod,fct[i]=1llu*i*fct[i-1]%mod;
    inv[n]=qpow(fct[n],mod-2);
    for(int i=n-1;i>=0;i--)inv[i]=1llu*inv[i+1]*(i+1)%mod;
    dfs1(1,0);dfs2(1,0);
    for(int i=1;i<=n;i++)add(ans,g[i]);
    printf("%d\n",ans);
    return 0;
}

其实,我骗了你,并不是所有人都喜欢饼干。但又能怎么样呢,回家找妈妈哭吗?