Tree with Maximum Cost CodeForces - 1092F（換根dp）

You are given a tree consisting exactly of n vertices. Tree is a connected undirected graph with n−1 edges. Each vertex v of this tree has a value av assigned to it.

Let dist(x,y) be the distance between the vertices x and y. The distance between the vertices is the number of edges on the simple path between them.

Let’s define the cost of the tree as the following value: firstly, let’s fix some vertex of the tree. Let it be v. Then the cost of the tree is ∑i=1ndist(i,v)⋅ai.

Your task is to calculate the maximum possible cost of the tree if you can choose v arbitrarily.

Input
The first line contains one integer n, the number of vertices in the tree (1≤n≤2⋅105).

The second line of the input contains n integers a1,a2,…,an (1≤ai≤2⋅105), where ai is the value of the vertex i.

Each of the next n−1 lines describes an edge of the tree. Edge i is denoted by two integers ui and vi, the labels of vertices it connects (1≤ui,vi≤n, ui≠vi).

It is guaranteed that the given edges form a tree.

Output
Print one integer — the maximum possible cost of the tree if you can choose any vertex as v.

Examples
Input
8
9 4 1 7 10 1 6 5
1 2
2 3
1 4
1 5
5 6
5 7
5 8
Output
121
Input
1
1337
Output
0
Note
Picture corresponding to the first example:

You can choose the vertex 3 as a root, then the answer will be 2⋅9+1⋅4+0⋅1+3⋅7+3⋅10+4⋅1+4⋅6+4⋅5=18+4+0+21+30+4+24+20=121.

In the second example tree consists only of one vertex so the answer is always 0.

題意：
求每個點到其他點的距離乘以點權的和。

思路：
換根dp裸題。

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <cmath>
#include <vector>

using namespace std;
typedef long long ll;
const int maxn = 4e5 + 7;

int siz[maxn],a[maxn];
int head[maxn],nex[maxn * 2],to[maxn * 2],tot;
ll f[maxn],sum[maxn];
int n;

void add(int x,int y) {
    to[++tot] = y;
    nex[tot] = head[x];
    head[x] = tot;
}

void dfs(int u,int fa) {
    siz[u] = 1;
    sum[u] = a[u];
    for(int i = head[u];i;i = nex[i]) {
        int v = to[i];
        if(v == fa) continue;
        dfs(v,u);
        siz[u] += siz[v];
        sum[u] += sum[v];
        f[u] += f[v] + sum[v];
    }
}

void dfs2(int u,int fa) {
    for(int i = head[u];i;i = nex[i]) {
        int v = to[i];
        if(v == fa) continue;
        ll num = f[u] - sum[v] - f[v];
        f[v] += (sum[1] - sum[v]) + num;
        dfs2(v,u);
    }
}

int main() {
    scanf("%d",&n);
    for(int i = 1;i <= n;i++) {
        scanf("%d",&a[i]);
    }
    for(int i = 1;i < n;i++) {
        int x,y;scanf("%d%d",&x,&y);
        add(x,y);add(y,x);
    }
    dfs(1,-1);
    dfs2(1,-1);
    ll ans = 0;
    for(int i = 1;i <= n;i++) {
        ans = max(ans,f[i]);
    }
    printf("%lld\n",ans);
    return 0;
}