CF1092F Tree with Maximum Cost

题面

给你一个 (n) 个节点的树，每个点都有一个权值 (a_i)

定义 (dist(u, v)) 为 (u) 和 (v) 两点之间的权值。

找到一个 (u)，使得：

[sum_{i = 1}^{n} dist(i, u) a_i ]

最大，求出这个最大值。

solution

换根 (dp)

任取一个点 (u) 当做根。

那么它的价值就是 (sum_{i = 1}^{i = n} a_i imes dep_i)

一遍 (dfs) ，处理处以 (u) 点为根的所有节点的子树内的价值。

(siz_u) 是子树内节点的权值和。

(val[u] = sum val[v] + siz[v])

因为多了 (u ightarrow v) 所以 (v) 子树内的节点深度都 (+1)

这样就求出了以 (u) 为根的答案 (val[u])

然后换根计算贡献：

(val[v] = sum + (val[u] - val[v]) - val[v])

code

#include<bits/stdc++.h>
#define int long long
using namespace std;
const int MAXA = 1e4 + 5;
const int MAXB = 2e5 + 5;
const int MAXC = 1e6 + 5;
int read() {
  int x = 0, f = 1; char c = getchar();
  while(c < '0' || c > '9') {if(c == '-') f = -1;c = getchar();}
  while(c >= '0' && c <= '9') {x = x * 10 + c - '0';c = getchar();}
  return x * f;
}
int n, a[MAXB], val[MAXB], siz[MAXB], f[MAXB];
struct edge{int v, nxt;}e[MAXB << 1];
int E, head[MAXB];
void add_edge(int u, int v) {
  e[++E] = (edge){v, head[u]};
  head[u] = E;
}
void dfs(int x, int fa) {
   siz[x] = a[x], val[x] = 0;
   for (int i = head[x]; i; i = e[i].nxt) {
   	  int v = e[i].v;
   	  if(v == fa) continue;
   	  dfs(v, x);
   	  siz[x] += siz[v];
	  val[x] += val[v] + siz[v];
   }
}
void dfs2(int x, int fa, int sum) {
   for (int i = head[x]; i; i = e[i].nxt) {
   	    int v = e[i].v;
   	    if(v == fa) continue;
   	    f[v] = f[x] + sum + siz[x] - siz[v] - siz[v];
   	    dfs2(v, x, sum + siz[x] - siz[v]);
   }
}
signed main() {
   n = read();
   for (int i = 1; i <= n; i++) a[i] = read();
   for (int i = 1, u, v; i < n; i++) {
   	  u = read(), v = read();
   	  add_edge(u, v), add_edge(v, u);
   }
   dfs(1, 0);
   f[1] = val[1];
   dfs2(1, 0, 0);
   int Ans = -0x3f3f3f3f3f3f;
   for (int i = 1; i <= n; i++) Ans = max(Ans, f[i]);
   cout<<Ans;
   return 0;
}