C++14(g++5.4) 解法, 执行用时: 261ms, 内存消耗: 22240K, 提交时间: 2019-12-27 19:15:52

#include <bits/stdc++.h>
#define ll long long
using namespace std;
const int N=2e5+5;
struct P {
	int v,w,id;
};
vector<P> G[N];
int a[N],num[N],n;
ll S,sum[N],ans[N];
void dfs(int u,int fa) {
	sum[u]=a[u],num[u]=1;
	for(int i=0;i<G[u].size();i++) {
		int v=G[u][i].v,w=G[u][i].w,id=G[u][i].id;
		if(v==fa) continue;
		dfs(v,u);
		sum[u]+=sum[v];
		num[u]+=num[v];
		ans[id]=(sum[v]*(n-num[v])+(S-sum[v])*num[v])*w;
	}
}
int main() {
	scanf("%d",&n);
	for(int i=1;i<=n;i++) {
		scanf("%d",&a[i]);
		S+=a[i];
	}
	for(int i=1,u,v,w;i<n;i++) {
		scanf("%d%d%d",&u,&v,&w);
		G[u].push_back({v,w,i});
		G[v].push_back({u,w,i});
	}
	dfs(1,-1);
	for(int i=1;i<n;i++) {
		printf("%lld\n",ans[i]*2);
	}
	return 0;
}

C++11(clang++ 3.9) 解法, 执行用时: 664ms, 内存消耗: 22928K, 提交时间: 2020-01-02 20:56:47

#include<bits/stdc++.h>
#define ll long long
using namespace std;
const int n=2e5+5;
struct node
{
	int v,w,id;
};
vector<node>G[n];
int a[n],s,r[n],h[n],t;
ll ans[n];
void dfs(int u,int aft)
{
	r[u]=a[u];h[u]=1;
	for(int i=0;i<G[u].size();i++)
	{
	   int v=G[u][i].v;int id=G[u][i].id;int w=G[u][i].w;
		 if(v==aft) continue;
		 dfs(v,u);
		 r[u]+=r[v];
		 h[u]+=h[v];
		 ans[id]=((1ll*((s-r[v])*h[v]))+(1ll*(t-h[v])*r[v]))*w*2;
	}
}
int main()
{
	cin>>t;
	for(int i=1;i<=t;i++){cin>>a[i];s+=a[i];}
	for(int i=1;i<t;i++)
	{
		int u,v,w;
		cin>>u>>v>>w;
		G[u].push_back({v,w,i});
		G[v].push_back({u,w,i});
	}
	dfs(1,-1);
	for(int i=1;i<t;i++)
	cout<<ans[i]<<'\n';
}