[CF1092F] Tree with Maximum Cost

Description

有 (n) 个节点的树，每个节点有一点权 (a_i)。定义 ( extrm{dist}(x,y)) 为 (x) 到 (y) 的边数。选取一点 (v)，使 $ sum_{i=1}^{n}a_i cdot extrm{dist}(i ,v)$ 最大

Solution

典型换根 dp，设 (1) 号点为根，设 (s[i]) 表示点 (i) 子树中的权值和，(f[i]) 表示 (v=i) 时的答案，首先我们可以通过一次 DFS 算出 (f[1])，同时有转移方程

[f[q]=f[p]-s[q]+s[1]-s[q]=f[p]+s[1]-2s[q] ]

#include <bits/stdc++.h>
using namespace std;

#define int long long
const int N = 1000005;

vector <int> g[N];
int n,a[N],dep[N],f[N],s[N],vis[N];

void dfs1(int p) {
    vis[p]=1;
    f[1]+=a[p]*dep[p];
    s[p]=a[p];
    for(int q:g[p]) if(!vis[q]) {
        dep[q]=dep[p]+1;
        dfs1(q);
        s[p]+=s[q];
    }
}

void dfs2(int p) {
    vis[p]=1;
    for(int q:g[p]) if(!vis[q]) {
        f[q]=f[p]+s[1]-2*s[q];
        dfs2(q);
    }
}

signed main() {
    ios::sync_with_stdio(false);
    cin>>n;
    for(int i=1;i<=n;i++) cin>>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);
    }
    dfs1(1);
    memset(vis,0,sizeof vis);
    dfs2(1);
    cout<<*max_element(f+1,f+n+1);
}