B 小琛和他的学校

链接：https://ac.nowcoder.com/acm/contest/4047/B
来源：牛客网

题目描述
小琛是一所学校的校长。
他的学校有n个校区（编号1~n），被n-1条双向道路连接，呈树形结构。
第i个校区共有Ai个学生。
第i天早上，所有的学生会沿最短路走到第i个校区参加活动，晚上再原路返回。
一个人通过第j条通道一次（即一人次），需要小琛支付wj的维护费用。
小琛想知道第n天结束之后，对于每一条通道，他总共需要支付多少费用。
对于100%的数据，1≤ n ≤ 200,000，1≤ A[i]≤ 10,000，1≤ w[i] ≤ 10,000。
输入描述:

第一行一个整数n，表示校区的数量。
接下来一行，n个整数，表示A1~An。
第3到第n+1行，每行包含3个整数。第i行包含三个整数ui-2,vi-2,wi-2，表示第i-2条通道所连接的两个校区的编号，以及一人次通过这条通道的费用。

输出描述:

共n-1行，每行一个整数。
第i行的整数表示小琛对于第i条通道所需支付的费用。

示例1
输入
复制

4
2 1 2 3
1 3 1
1 2 3
4 1 2

输出
复制

24
60
56

利用dfs的回溯来进行一些操作，和Tarjan有点类似。

#include <iostream>
#include <algorithm>
#include <cstdio>
#include <cstring>
using namespace std;

const int maxn = 2e5 + 10;
typedef long long ll;
struct node{
	int u, v, nex, id, cost;
}edge[maxn * 2];
ll n, ans[maxn];
int head[maxn];
ll a[maxn], cnt, tot, maxSum;
ll sz[maxn], sum[maxn];

void addedge(int u, int v, int id, int cost) {
	edge[cnt] = {u, v, head[u], id, cost};
	head[u] = cnt++;
}

void init() {
	memset(head, -1, sizeof(head));
}

void dfs(int u, int fa) {
	sz[u] = 1;
	sum[u] = a[u];
	for(int i = head[u]; i != -1; i = edge[i].nex) {
		int v = edge[i].v;
		if(v == fa)
			continue;
		dfs(v, u);
		sum[u] += sum[v];
		sz[u] += sz[v];
		int id = edge[i].id;
		ans[id] = (sz[v] * (maxSum - sum[v]) + (n - sz[v]) * sum[v]) * edge[i].cost;
	}
}

int main() {
	init();
	ios::sync_with_stdio(false);
	cin >> n;
	for(int i = 1; i <= n; i++) {
		cin >> a[i];
		maxSum += a[i];
	}
	for(int i = 1; i < n; i++) {
		int u, v, data;
		cin >> u >> v >> data;
		addedge(u, v, i, data);
		addedge(v, u, i, data);
	}
	dfs(1, 0);
	for(int i = 1; i < n; i++)
		cout << ans[i] * 2 << endl;
	return 0;
}