fork download 
  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 sz[MAXN+1];        // subtree sizes
  11. ll distSum[MAXN+1];   // distSum[u] = sum of dist(u, i) over all i
  12. ll ans = LLONG_MIN;
  13.  
  14. // First DFS: compute sz[u] and distSum[1]
  15. ll totalDist1 = 0;
  16. void dfs1(int u, int p, int depth) {
  17.     sz[u] = 1;
  18.     totalDist1 += depth;
  19.     for (int v : adj[u]) {
  20.         if (v == p) continue;
  21.         dfs1(v, u, depth+1);
  22.         sz[u] += sz[v];
  23.     }
  24. }
  25.  
  26. // Second DFS: reroot to compute distSum[u] for all u
  27. void dfs2(int u, int p) {
  28.     // record answer for u
  29.     ans = max(ans, a[u] * distSum[u]);
  30.     for (int v : adj[u]) {
  31.         if (v == p) continue;
  32.         // when moving root from u -> v:
  33.         // distSum[v] = distSum[u] - sz[v] + (n - sz[v])
  34.         distSum[v] = distSum[u] - sz[v] + (n - sz[v]);
  35.         dfs2(v, u);
  36.     }
  37. }
  38.  
  39. int main(){
  40.     ios::sync_with_stdio(false);
  41.     cin.tie(nullptr);
  42.  
  43.     cin >> n;
  44.     for (int i = 1; i <= n; i++) {
  45.         cin >> a[i];
  46.     }
  47.     for (int i = 0; i < n-1; i++) {
  48.         int u, v;
  49.         cin >> u >> v;
  50.         adj[u].push_back(v);
  51.         adj[v].push_back(u);
  52.     }
  53.  
  54.     // 1) Compute subtree sizes and totalDist1 = sum distances from 1
  55.     dfs1(1, 0, 0);
  56.     distSum[1] = totalDist1;
  57.  
  58.     // 2) Reroot DP to fill distSum[u] for all u and track max cost
  59.     dfs2(1, 0);
  60.  
  61.     cout << ans << "\n";
  62.     return 0;
  63. }
  64.  
Success #stdin #stdout 0.01s 11876KB
comments (?)
stdin 
8
9 4 1 7 10 1 6 5
1 2
2 3
1 4
1 5
5 6
5 7
5 8
stdout 
119
https://ideone.com/8w7oir
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!