ICPC NEAU Programming Contest 2020—— G. 選根

有一顆有n個結點樹，結點被編號爲1 ~ n，記根結點深度爲1，如果第iii個結點的深度是d，則它貢獻的價值是d×wi，這棵樹的價值是所有結點的價值和 求當根結點爲1~n時，樹的價值分別爲多少

輸入描述

第一行輸入一個整數T，代表有T組測試數據 對於每一組測試數據，第一行有1個整數n，第二行有n個整數wi​，接下來n−1行每行有兩個整數x,y表示x和y之間有一條邊

輸出描述

對於每組測試數據，在一行中輸出n個整數，第i個整數代表以i結點爲根結點時樹的價值

數據範圍

1≤T≤1000
1≤n≤2⋅10⁵
1≤wi≤10⁸
∑n≤2⋅10⁵

輸出時每行末尾的多餘空格，不影響答案正確性
樣例輸入

2
6
5 2 8 1 7 8
4 5
5 6
2 5
1 3
4 3
5
1 1 1 1 1
1 2
2 3
3 4
4 5

樣例輸出

102 100 81 76 73 88
15 12 11 12 15

首先定義數組sum[x]（以x爲根的樹的權值和）、ans[x]（以x爲根的樹的價值），先以1爲根算出樹的價值ans[1]，它兒子j爲根的樹的價值ans[j] = ans[1] + sum[1] - 2*sum[j]，依次類推即可。

#include<bits/stdc++.h>

using namespace std;

const int maxn = 1e6+1;
long long w[maxn], sum[maxn], ans[maxn];

struct edge{
    int to, next;
    edge(){}
    edge(int to, int next): to(to), next(next){}
}e[maxn]; int head[maxn];

void dfs1(int x, int pre, int d) {
    sum[x] = w[x];
    ans[x] = w[x] * d;
    for(int i = head[x]; ~i; i = e[i].next) {
        int j = e[i].to;
        if (j == pre) continue;
        dfs1(j, x, d+1);
        sum[x] += sum[j];
        ans[x] += ans[j];
    }
}

void dfs2(int x, int pre) {
    for(int i = head[x]; ~i; i = e[i].next) {
        int j = e[i].to;
        if (j == pre) continue;
        ans[j] = ans[x] + sum[x] - 2*sum[j];
        sum[j] = sum[x];
        dfs2(j, x);
    }
}

int main() {   
    int T;
    scanf("%d", &T);
    while(T --) { 
        memset(head, -1, sizeof(head));
        int n, tot = 0;
        scanf("%d", &n);
        for(int i = 1; i <= n; ++ i) {
            scanf("%lld", &w[i]);
        }
        for(int i = 1; i < n; ++ i) {
            int u, v;
            scanf("%d%d", &u, &v);
            e[++tot] = edge(u, head[v]), head[v] = tot;
            e[++tot] = edge(v, head[u]), head[u] = tot;
        }
        dfs1(1, 1, 1);
        dfs2(1, 1);
        for(int i = 1; i <= n; ++ i) {
            printf("%lld ", ans[i]);
        }
        printf("\n");
    }
    return 0;
}