Radio Prize（换根dp）

题意

给定一棵无根树，树上每个节点有点权$t_i$，两点之间有边权$w_i$。定义从$u$到$v$的距离$d_{u,v}$为从$u$到$v$路径的边权和，花费为$(t_u+t_v)\ast d_{u,v}$。对每个点，都要求出到其他点花费的总和。

数据范围

$1\le n\le 100000$

思路

首先先把花费的式子拆开，即$p_{u,v}=(t_u+t_v)\ast d_{u,v}=t_u\ast d_{u,v}+t_v\ast d_{u,v}$。假设要求$u$到其他点花费的总和，那么$P_u=t_u\ast {\sum }_v{d}_{u,v}+{\sum }_v(t_v\ast d_{u,v})$。令$dis{t}_u={\sum }_v{d}_{u,v}$，$distva{l}_u={\sum }_v(t_v\ast d_{u,v})$，则$P_u=t_u\ast dis{t}_u+distva{l}_u$。
我们再探讨$dis{t}_u$以及$distva{l}_u$的实际意义，$dis{t}_u$表示其他点到$u$点的距离总和，$distva{l}_u$表示其他点到$u$点的距离与点权乘积之和。我们可以先预处理出以$u$为根的子树中节点到$u$点的距离总和以及以$u$为根的子树中节点到$u$点的距离与点权成绩之和。假设父节点为$x$，子节点为$y$，则预处理的转移方程为$dist[x]=dist[y]+sumn[y]\ast w[i]$，这里的$sumn$数组指的是子树中节点的数量。$distval[x]=distval[y]+sumv[y]\ast w[i]$，这里的$sumv$数组指的是子树中节点点权之和。
然后我们再通过换根dp的方式，转移出$dist,distval$数组。假设父节点为$x$，子节点为$y$，转移方程为$dist[y]=dist[x]-sumn[y]\ast w[i]+(n-sumn[y])\ast w[i]$，这里的$n$指的是节点的数量。$distval[y]=distval[x]-sumv[y]\ast w[i]+(V-sumv[y])\ast w[i]$，这里的$V$指的是所有节点的点权总和。
理顺一下思路，就是先预处理出$sumn,sumv,dist,distval$四个数组，然后再通过树形dp的方式更新$dist,distval$两个数组，最后再通过推出的公式计算出答案。

代码

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>

using namespace std;

const int N = 100010, M = 2 * N;

typedef long long ll;

int n;
ll val[N];
int h[N], e[M], ne[M], idx;
ll w[M];
ll sumn[N], sumv[N];
ll dist[N], distval[N];
ll sumvalue;

void add(int a,int b,ll c)
{
    e[idx] = b, w[idx] = c, ne[idx] = h[a], h[a] = idx ++;
}

void dfs1(int u,int fa)
{
    for(int i=h[u];~i;i=ne[i]){
        int j = e[i];
        if(j == fa) continue;
        dfs1(j,u);
        sumn[u] += sumn[j];
        sumv[u] += sumv[j];
        dist[u] += dist[j] + sumn[j] * w[i];
        distval[u] += distval[j] + sumv[j] * w[i];
    }
    sumn[u] ++;
    sumv[u] += val[u];
}

void dfs2(int u,int fa)
{
    for(int i=h[u];~i;i=ne[i]){
        int j = e[i];
        if(j == fa) continue;
        dist[j] = dist[u] - w[i] * sumn[j] + (n - sumn[j]) * w[i];
        distval[j] = distval[u] - w[i] * sumv[j] + (sumvalue - sumv[j]) * w[i];
        dfs2(j,u);
    }
}

int main()
{
    cin >> n;
    memset(h,-1,sizeof(h));
    for(int i=1;i<=n;i++){
        cin >> val[i];
        sumvalue += val[i];
    }
    for(int i=1;i<n;i++){
        int a,b,c;
        cin >> a >> b >> c;
        add(a,b,c), add(b,a,c);
    }
    dfs1(1,-1);
    dfs2(1,-1);
    for(int i=1;i<=n;i++) cout << distval[i] + val[i] * dist[i] << "\n";
    return 0;
}