fork download
copy 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3. using ll = long long;
  4.  
  5. const int MAXN = 200000;
  6.  
  7. int n;
  8. ll a[MAXN+1];
  9. vector<int> adj[MAXN+1];
  10. ll sumA[MAXN+1];      // sum of a[] in subtree
  11. ll distSum[MAXN+1];   // weighted distance sum at each root
  12. ll totalA = 0;
  13. ll answer = LLONG_MIN;
  14.  
  15. // First DFS: compute sumA[u] and distSum[1] (rooted at 1)
  16. void dfs1(int u, int p, int depth) {
  17.     sumA[u] = a[u];
  18.     distSum[1] += a[u] * depth;
  19.     for (int v : adj[u]) {
  20.         if (v == p) continue;
  21.         dfs1(v, u, depth + 1);
  22.         sumA[u] += sumA[v];
  23.     }
  24. }
  25.  
  26. // Second DFS: reroot from u to v
  27. void dfs2(int u, int p) {
  28.     // update answer at u
  29.     answer = max(answer, distSum[u]);
  30.     for (int v : adj[u]) {
  31.         if (v == p) continue;
  32.         // when rerooting from u -> v:
  33.         // nodes in v-subtree get distance -1 each (sumA[v])
  34.         // nodes outside v-subtree get distance +1 each (totalA - sumA[v])
  35.         distSum[v] = distSum[u] + totalA - 2*sumA[v];
  36.         dfs2(v, u);
  37.     }
  38. }
  39.  
  40. int main(){
  41.     ios::sync_with_stdio(false);
  42.     cin.tie(nullptr);
  43.  
  44.     cin >> n;
  45.     for (int i = 1; i <= n; i++) {
  46.         cin >> a[i];
  47.         totalA += a[i];
  48.     }
  49.     for (int i = 0; i < n-1; i++) {
  50.         int u, v;
  51.         cin >> u >> v;
  52.         adj[u].push_back(v);
  53.         adj[v].push_back(u);
  54.     }
  55.  
  56.     // 1) dfs1 computes sumA[] and distSum[1]
  57.     dfs1(1, 0, 0);
  58.  
  59.     // 2) dfs2 reroots and fills distSum[u] for all u
  60.     dfs2(1, 0);
  61.  
  62.     // Each cost is distSum[u], but original problem multiplies by a_v?
  63.     // Actually cost is sum_i dist(i,v)*a_i = distSum[v].
  64.     // If the problem wanted a_v * that, replace answer = max(answer, a[u]*distSum[u]);
  65.     cout << answer << "\n";
  66.     return 0;
  67. }
  68.  
Success #stdin #stdout 0.01s 11692KB
comments (?)
stdin
copy 
8
9 4 1 7 10 1 6 5
1 2
2 3
1 4
1 5
5 6
5 7
5 8
stdout
copy 
121
https://ideone.com/PFvaa6
language:
C++ (gcc 8.3)
created:
3 days ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!