Radio Prize 思维/树上dp

题目链接

题意

给定一棵树，树上每条边及每个点都有权值。定义$P_u={\sum }_{v}d(u,v)(t_u+t_v)$其中$d(u,v)$为树上u到v的距离，$t_u$和$t_v$分别是两点权值，求每个点的P值。

分析

首先我们可以看到$P_u$可以分成 $t_u{\sum }_{v}d(u,v)+{\sum }_v{t}_{v}d(u,v)$ ，则我们只需要知道对于每一个节点${\sum }_{v}d(u,v)$和${\sum }_v{t}_{v}d(u,v)$的值即可。直接计算复杂度太高，我们可以考虑通过一个已经计算出值的节点计算当前节点值。

当我们已经计算得到$u$的值，并想要得到与$u$相邻$u^{\prime }$的值时，可以将除了$u^{\prime }$的其他节点分为两部分，一部分不经过$u$，另一部分经过$u$。

• 计算${\sum }_{v}d(u,v)$时，对于这两部分节点来说，不经过$u$的节点的$d(u^{\prime },v)$相比$d(u,v)$减少了$w_i$，经过$u$的节点的$d(u^{\prime },v)$相比$d(u,v)$增加了$w_i$，故只需知道两部分节点的分别数量即可得到${\sum }_{v}d(u,v)$。

• 计算${\sum }_v{t}_{v}d(u,v)$时，对于这两部分节点来说，不经过$u$的节点的$t_{v}d(u^{\prime },v)$相比$t_{v}d(u,v)$减少了$w_i{\sum }_{v}tv$，经过$u$的节点的$t_{v}d(u^{\prime },v)$相比$t_{v}d(u,v)$增加了了$w_i{\sum }_{v}tv$，所以这部分只需知道两部分节点的点值和即可得到${\sum }_v{t}_{v}d(u,v)$

有了转移方法后，我们需要获得对于每一个节点两部分分别的值，故可以指定任意节点为根节点，通过一次dfs，得到子节点数量和及子节点点值和后，通过所有节点数量和所有节点点值和得到另一部分节点的对应值，即可进行转移。

代码

#include <bits/stdc++.h>

using namespace std;
typedef long long ll;
const int N = 5e5 + 10;
long long sum1[N],sum2[N],num[N],sum[N],Sum;
struct node{
    int u,v,w,next;
}a[N];
int n;
int f[N],cnt;
long long value[N],dis[N];
void add(int u,int v,int w){
    a[++cnt].v =v;
    a[cnt].w = w;
    a[cnt].next = f[u];
    f[u] = cnt;
}
int vis[N];

void dfs1(int u){
    num[u] = 1;
    sum[u] = value[u];
    for(int i = f[u];i;i = a[i].next){
        int v = a[i].v,w = a[i].w;
        if(!vis[v]){
            vis[v] = 1;
            dis[v] = dis[u]+w;
            sum1[1] += dis[v];
            sum2[1] += dis[v]*value[v];
            dfs1(v);
            num[u] += num[v];
            sum[u] += sum[v];
        }
    }
}
long long ans[N];
void dfs2(int u){
    ans[u] = sum1[u]*value[u] + sum2[u];
    for(int i = f[u];i;i = a[i].next){
        int v = a[i].v,w = a[i].w;
        if(!vis[v]){
            vis[v] = 1;
            sum1[v] = sum1[u] + (n-num[v]*2)*w;
            sum2[v] = sum2[u] + (Sum-sum[v]*2)*w;
            dfs2(v);
        }
    }
}
int main() {
    cin>>n;
    for(int i = 1;i<=n;i++){
        scanf("%d",&value[i]);
        Sum += value[i];
    }
    int u,v,w;
    for(int i = 1;i<n;i++){
        scanf("%d %d %d",&u,&v,&w);
        add(u,v,w);
        add(v,u,w);
    }
    memset(vis,0,sizeof vis);
    vis[1] = 1;
    dfs1(1);
    memset(vis,0,sizeof vis);
    vis[1] = 1;
    dfs2(1);
    for(int i = 1;i<=n;i++){
        printf("%lld\n",ans[i]);
    }
    return 0;
}