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;
}