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