Problem:

1092F - Tree with Maximum Cost

Problem description:

Given a tree with n vertices, where each vertex has a number a_i written on it. Distance between two vertices is defined as the number of edges on the path between them. You have to choose a vertex v such that d₁ × a₁ + d₂ × a₂ + d₃ × a₃ + ... ... + d_n × a_n is maximized, where d_i denotes the distance between vertex i and vertex v. You have to print this maximum value.

Solution:

Lets tot be the sum of all a_i for all vertices i and sum_i be the sum of all a_j of the vertices j in the subtree of vertex i.

Now, lets subtot_u stores the sum of all d_i × a_i for all vertices i in the subtree of vertex u, where d_i is the distance of vertex i from the vertex u. But how to calculate the array subtot?

Lets root the tree at vertex 1. If u is a leaf vertex, then it is obvious that subtot[u] = 0 × a[u] = 0, since distance from u tot u is 0.

Now what about the non leaf nodes?

For a vertex v, subtot[v] stores all a[i] × d[i] for all vertex i in the subtree of v. Since for its parent u, it is 1 distance upper than it's child v, we need to add sum[v] with subtot[v] to get subtot[u], for all child v of u.

So, for any non leaf vertex u, subtot[u] = subtot[v] + sum[v], for all direct child v of vertex u. Addition of sum[v] is for 1 distance away from child v to parent u.

So, run a dfs to calculate the array sum[] and subtot[].

vector<int> g[maxn];
ll subtot[maxn], sum[maxn];
void dfs(int u, int p){
    subtot[u] = 0LL;
    sum[u] = a[u];
    for(int v:g[u]) if(v!=p){
        dfs(v,u);
        sum[u]+=sum[v];
        subtot[u] += subtot[v] + sum[v];
    }
}

Ok, now if we choose vertex u as the fixed vertex, how to calculate the answer for this vertex?

Notice that, the calculation for any vertex which is in the subtree of vertex u will not change. What will be changed? Other vertices which is not in the subtree of vertex u. Lets the sum of all subtot[] values of other vertices which are not in the subtree of vertex u be up. Every time you go one step down, distance from vertex u to all other vertices which are not in the subtree of vertex u will increase by 1. So, wee need to increase up by tot - sum[u](we are ignoring the sum[u], since distance to other nodes is increased, which is, sum of all a[i], for all vertex i which are not in the subtree of vertex u.)

Then the current answer for vertex u is cur = subtot[u] + up + tot - sum[u]. Here, subtot[u] is for for vertices in the subtree of u, up for vertices which are not in the subtree, tot-sum[u] is for increasing ditance by 1 for non subtree vertices.

When you will go to vertex v from vertex u, then update up by up = cur - subtot[v] - sum[v]. Here since we are going down, that's why we subtract subtot[v] and sum[v], instead of adding.

Thus, you have to run another dfs to try to choose each vertex as the fixed vertex and keep track of the maximum value cur ans our required answer in ans.

ll ans = 0LL;
void dfs2(int u, int p, ll up){
    ll cur = subtot[u] + up + tot-sum[u];
    ans = max(ans, cur);
    for(int v:g[u]) if(v!=p){
        dfs2(v, u, cur - subtot[v] - sum[v]);
    }
}

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int maxn = 200002;
int n;
ll a[maxn], tot, ans;
vector<int> g[maxn];
ll subtot[maxn], sum[maxn];
void dfs(int u, int p){
    subtot[u] = 0LL;
    sum[u] = a[u];
    for(int v:g[u]) if(v!=p){
        dfs(v,u);
        sum[u]+=sum[v];
        subtot[u] += subtot[v] + sum[v];
    }
}
void dfs2(int u, int p, ll up){
    ll cur = subtot[u] + up + tot-sum[u];
    ans = max(ans, cur);
    for(int v:g[u]) if(v!=p){
        dfs2(v, u, cur - subtot[v] - sum[v]);
    }
}
int main()
{
    ios_base::sync_with_stdio(0); cin.tie(NULL); cout.tie(NULL);
    cin>>n;
    for(int i=1;i<=n;i++) {
        cin>>a[i];
        tot+=a[i];
    }
    for(int i=1;i<n;i++){
        int u,v; 
        cin>>u>>v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    dfs(1,0);
    dfs2(1,0,0LL);
    cout << ans;
    return 0;
}

Here is my submission link: 47272565

If you have any further query regarding any part of the solution, feel free to ask in the comment.