树形DP的一些题

问题1

(http://zhengruioi.com/problem/1030)
\(n\) 个点的树，点有点权，定义一个连通块的贡献为其中所有点的点权和的平方。现在要求所有连通块的贡献之和。\(n\leq 5\times 10^5\)

sol：

考虑把平方式展开，\((a+b)^2=a^2+b^2+2ab\) 那么实质上就是要维护三个数组\(f[x], g[x], h[x]\)，分别表示\(x\) 子树所有连通块权值和的平方，权值的和，还有连通块的个数。转移就很简单了。

#include <bits/stdc++.h>
#define ll long long
using namespace std;
inline int read(){
    int x=0;char c=getchar();
    while(c<'0'||c>'9') c=getchar();
    while(c>='0'&&c<='9') x=(x<<1)+(x<<3)+c-'0',c=getchar();
    return x;
}
const int N=500005,mod=998244353;
ll f[N],g[N],h[N],ans;
int n,w[N],ver[N<<1],nxt[N<<1],head[N],tot;
inline void link(int x,int y){ver[++tot]=y;nxt[tot]=head[x];head[x]=tot;}
void dfs(int x,int la){
    h[x]=1;g[x]=w[x];f[x]=1ll*w[x]*w[x]%mod;
    for(int i=head[x];i;i=nxt[i]){
        int y=ver[i];
        if(y==la) continue;
        dfs(y,x);
        f[x]=(f[x]+f[x]*h[y]%mod+g[x]*g[y]*2%mod+f[y]*h[x]%mod)%mod;
        g[x]=(g[x]+g[x]*h[y]%mod+g[y]*h[x]%mod)%mod;
        h[x]=(h[x]+h[x]*h[y]%mod)%mod;
    }
    ans=(ans+f[x])%mod;
}
int main(){
    n=read();
    for(int i=1;i<=n;i++) w[i]=read();
    for(int i=1;i<n;i++){
        int x=read(),y=read();
        link(x,y);link(y,x);
    }
    dfs(1,0);
    printf("%d\n",ans);
    return 0;
}

问题2

\(n\) 个点的树，点有点权，定义一条简单路径的贡献为其中所有点的点权和的平方。现在要求所有简单路径的贡献之和。\(n\leq 5\times 10^5\)

sol：

还是维护这些东西，但转移就相对没有那么好写了。其实道理是一样的。

问题3

(http://zhengruioi.com/contest/232/problem/595)
\(n\) 个点的树，边有边权，现在要按一个排列访问所有节点，从一个点\(u\) 到\(v\) 消耗的体力为\(u\) 到\(v\) 的简单路径上的所有边的长度的乘积，设一个方案一共消耗的体力为每次从一个点到另一个点消耗的体力之和。现在要求所有排列耗费体力之和。\(n\leq 5\times 10^5\)

sol：

真的不敢想象曾经普及组的我竟在模拟赛上切了这一题。因为是全排列，所以每一个点对\((u,v)\) 都会被算\((n-1)!\times 2\) 遍。考虑如何求所有\((u,v)\) 之和。维护\(f[x], g[x]\) 分别表示\(x\) 子树内使\(x\) 度为\(1\) 的路径，使\(x\) 度为\(2\) 的路径。转移见代码。

#include <bits/stdc++.h>
#define ll long long
using namespace std;
const int N=500005,mod=998244353;
int n,ver[N<<1],nxt[N<<1],head[N],tot=1;
ll f[N],edge[N<<1],g[N];
void add(int x,int y,int z){ver[++tot]=y;edge[tot]=z;nxt[tot]=head[x];head[x]=tot;}
void dp(int x,int la){
    for(int i=head[x];i;i=nxt[i]){
        int y=ver[i];
        if(y==la) continue;
        dp(y,x);
        g[x]=(g[x]+edge[i]*f[y]%mod*f[x])%mod;
        f[x]=(f[x]+edge[i]*f[y]%mod)%mod;
    }
    f[x]++;
}
int main(){
    scanf("%d",&n);
    for(int i=1;i<n;i++){
        int x,y,z;scanf("%d%d%d",&x,&y,&z);
        add(x,y,z);add(y,x,z);
    }
    dp(1,0);
    long long ans=0,step=1;
    for(int i=1;i<=n;i++) ans=(ans+g[i]+f[i]-1)%mod;
    for(int i=1;i<=n-1;i++) step=(step*i)%mod;
    printf("%lld\n",ans*step*2%mod);
    return 0;
}

CF1249F

(http://codeforces.com/contest/1249/problem/F)
\(n\) 个点的树和一个限制\(k\)，点有点权，现在要选出一个点集，满足任意两个点的距离严格大于\(k\)，并使得点权之和最大，求最大值。\(n,k\leq 200\)

sol：