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 \times dep_i\)

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

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

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

因为多了 \(u\rightarrow 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;
}