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