AT_dp_v Subtree 题解

2023-07-24 20:16:31 AT_dp_v 题解

AT_dp_v Subtree

思路

考虑树形 dp，假设根左右子树的连通块已经算出来了，我们用 $f[x]$ 表示强制将 $x$ 染色，$x$ 子树的连通块方案数。

那么我们可得：
$f[u]=\prod _{v=u_{son}}(f[v]+1)$

$+1$ 是因为还要算上这个子树不取的情况。

但是考虑到根的答案处理出来后，其他节点的答案还没有处理出来。由于 $n\le {10}^5$，我们显然不能对每个节点设根做一遍 dfs，所以考虑换根 dp。

由于在换根之后根节点除了这个子树以外的部分都算这个点的子树，转移方法同上，那我们思考后又可得:

$f^{\prime }[v]=f[v]\times (f[u]÷(f[v]+1)+1)$

但是我们又要对答案取模，这使得式子应该为：

$f^{\prime }[v]=f[v]\times (f[u]÷(f[v]+1)+1)modm$

而 $m$ 不一定为质数，这意味着我们不能通过求逆元的方式来解决这个除法。

那怎么办呢？

显然不能再算一遍，那我们只能想一个 $\log n$ 以下的方法处理。

考虑同时做前缀乘和后缀乘，计算对子树的贡献只需要将前缀乘和后缀乘相乘即可，虽然多了几个 $O(n)$，但是总复杂度不变。空间上，因为子树根的个数为 $O(n)$，所以空间复杂度也为 $O(n)$。

注意细节处理即可。

$Code:$

int n,mod,cnt;//cnt维护前缀乘起点。
ll dp[100005],pre[400005],sub[400005];//pre表示前缀乘数组，sub表示后缀乘数组，开4倍是为了放每个区间之间的2个间距
basic_string<int> f[100005];
void dfs1(int x,int fa){//处理根节点答案。
    dp[x]=1;
    for(int y:f[x]){
        if(y==fa)continue;
        dfs1(y,x);
        dp[x]=(dp[x]*(dp[y]+1))%mod;//套公式即可
    }
}
void dfs(int x,int fa,ll las){
    dp[x]=dp[x]*(las+1)%mod;
    int st=++cnt,en,o=0;//st表示开始位置，en表示结束位置（为什么要直接在数组上做？因为方便，而且数组快，内存有保证）
    cnt=cnt+f[x].size()+1;//这里直接加f[x].size()，不需要减一，因为父节点也要算进子节点的前缀后缀答案中（换根后作为一棵子树）
    pre[st-1]=1,sub[cnt-1]=1;//处理数组前后值。
    en=cnt-2;//cnt比en大2是为了放st-1和en+1。
    for(int i=0;i<f[x].size();i++){//前缀乘
        int y=f[x][i];
        if(y==fa)pre[st+o]=pre[st+o-1]*(las+1)%mod;//父亲加进来
        else pre[st+o]=pre[st+o-1]*(dp[y]+1)%mod;
        o++;
    }
    o=0;
    for(int i=f[x].size()-1;~i;--i){//后缀乘
        int y=f[x][i];
        if(y==fa)sub[en-o]=sub[en-o+1]*(las+1)%mod;
        else sub[en-o]=sub[en-o+1]*(dp[y]+1)%mod;
        o++;
    }
    for(int i=0;i<f[x].size();i++){
        int y=f[x][i];
        if(y==fa)continue;
        dfs(y,x,(pre[st+i-1]*sub[st+i+1])%mod);//递归
    }
}
int main(){
    n=read(),mod=read();
    for(int i=1;i<n;i++){
        int u=read(),v=read();
        f[u]+=v,f[v]+=u;
    }
    dfs1(1,0);
    dfs(1,0,0);
    for(int i=1;i<=n;i++){
        printf("%lld\n",dp[i]%mod);
    }
}