链接

https://ac.nowcoder.com/acm/contest/12606/I

题意

树上对于任意节点 u u u 求 ∑ v = 1 n [ d u , v ∗ ( t u + t v ) ] \sum_{v=1}^{n}{[d_{u,v}*(t_u+t_v)]} ∑v=1n​[du,v​∗(tu​+tv​)]

思路

换根 DP

记 r e s u = ∑ v = 1 n [ d u , v ∗ ( t u + t v ) ] res_u=\sum_{v=1}^{n}{[d_{u,v}*(t_u+t_v)]} resu​=∑v=1n​[du,v​∗(tu​+tv​)]

r e s u = ∑ v = 1 n [ d u , v ∗ ( t u + t v ) ] = t u ∗ ∑ v = 1 n d u , v + ∑ v = 1 n ( d u , v ∗ t v ) res_{u}=\sum_{v=1}^{n}{[d_{u,v}*(t_u+t_v)}]=t_u*\sum_{v=1}^{n}{d_{u,v}}+\sum_{v=1}^{n}{(d_{u,v}*t_v)} resu​=∑v=1n​[du,v​∗(tu​+tv​)]=tu​∗∑v=1n​du,v​+∑v=1n​(du,v​∗tv​)

记 a u = ∑ v = 1 n d u , v a_u=\sum_{v=1}^{n}{d_{u,v}} au​=∑v=1n​du,v​， b u = ∑ v = 1 n ( d u , v ∗ t v ) b_u=\sum_{v=1}^{n}{(d_{u,v}*t_v)} bu​=∑v=1n​(du,v​∗tv​)， s z u sz_u szu​ 为节点 u u u 的子树节点数， s u m u sum_u sumu​ 为节点 u u u 的子树节点权值和， u ′ u' u′ 为节点 u u u 的儿子节点， S U M = ∑ i = 1 n t i SUM=\sum_{i=1}^{n}{t_i} SUM=∑i=1n​ti​

a u ′ = a u − d u , u ′ ∗ s z u + d u , u ′ ∗ ( n − s z u ′ ) = a u + d u , u ′ ∗ ( n − 2 ∗ s z u ′ ) a_{u'}=a_u-d_{u,u'}*sz_u+d_{u,u'}*(n-sz_{u'})=a_u+d_{u,u'}*(n-2*sz_{u'}) au′​=au​−du,u′​∗szu​+du,u′​∗(n−szu′​)=au​+du,u′​∗(n−2∗szu′​)

b u ′ = b u − d u , u ′ ∗ s u m u + d u , u ′ ∗ ( S U M − s u m u ′ ) = b u + d u , u ′ ∗ ( S U M − 2 ∗ s u m u ′ ) b_{u'}=b_u-d_{u,u'}*sum_u+d_{u,u'}*(SUM-sum_{u'})=b_u+d_{u,u'}*(SUM-2*sum_{u'}) bu′​=bu​−du,u′​∗sumu​+du,u′​∗(SUM−sumu′​)=bu​+du,u′​∗(SUM−2∗sumu′​)

r e s u ′ = t u ′ ∗ a u ′ + b u ′ res_{u'}=t_{u'}*a_{u'}+b_{u'} resu′​=tu′​∗au′​+bu′​

两次 DFS， O ( n ) O(n) O(n) 求解

代码

#include <bits/stdc++.h>
#define SZ(x) (int)(x).size()
#define ALL(x) (x).begin(),(x).end()
#define PB push_back
#define EB emplace_back
#define MP make_pair
#define FI first
#define SE second
using namespace std;
typedef double DB;
typedef long double LD;
typedef long long LL;
typedef unsigned long long ULL;
typedef pair<int,int> PII;
typedef vector<int> VI;
typedef vector<PII> VPII;
// head
const int N=1e5+5;
int n;
LL t[N],d[N],sz[N],sum[N],a[N],b[N],res[N];
VPII g[N];
void dfs1(int u,int fa) {
    sz[u]=1;
    sum[u]=t[u];
    for(auto x:g[u]) {
        int v=x.FI,w=x.SE;
        if(v==fa) continue;
        d[v]=d[u]+w;
        a[1]+=d[v];
        b[1]+=d[v]*t[v];
        dfs1(v,u);
        sz[u]+=sz[v];
        sum[u]+=sum[v];
    }
}
void dfs2(int u,int fa) {
    for(auto x:g[u]) {
        int v=x.FI,w=x.SE;
        if(v==fa) continue;
        a[v]=a[u]+w*(n-2*sz[v]);
        b[v]=b[u]+w*(sum[1]-2*sum[v]);
        res[v]=a[v]*t[v]+b[v];
        dfs2(v,u);
    }
}
int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cin>>n;
    for(int i=1;i<=n;i++) cin>>t[i];
    for(int i=1;i<n;i++) {
        int u,v,w;
        cin>>u>>v>>w;
        g[u].EB(v,w);
        g[v].EB(u,w);
    }
    dfs1(1,0);
    dfs2(1,0);
    res[1]=a[1]*t[1]+b[1];
    for(int i=1;i<=n;i++) cout<<res[i]<<'\n';
    return 0;
}