fork download
copy 
  1. #include <bits/stdc++.h>
  2. #define ll long long
  3. #define pii pair<ll, ll>
  4. #define st first
  5. #define nd second
  6. #define file "test"
  7. using namespace std;
  8. const long long INF = 1e18;
  9. const long long N = 1e6 + 5;
  10.  
  11. ll n, w[N], dp[N], dp2[N], f[N], f2[N], g[N], ans = 0;
  12. vector <int> d[N];
  13.  
  14. void dfs1(int p, int u) {
  15.     dp[u] = f[u] = w[u];
  16.     vector <ll> save;
  17.     for (int v: d[u]) if (v ^ p) {
  18.         dfs1(u, v); 
  19.         dp[u] = max(dp[u], dp[v] + w[u]);
  20.         save.push_back(dp[v]);
  21.     }  
  22.  
  23.     sort(save.begin(), save.end(), greater <ll> ());
  24.     if (save.size() > 0) f[u] += save[0];
  25.     if (save.size() > 1) f[u] += save[1];
  26.  
  27.     for (int v: d[u]) if (v ^ p) f[u] = max(f[u], f[v]);
  28. }
  29.  
  30. void dfs2(int p, int u) {
  31.  
  32.     ans = max(ans, f[u]);
  33.  
  34.     vector <pii> save, ff;
  35.  
  36.     for (int v: d[u]) if (v != p) {
  37.         save.push_back({dp[v], v});
  38.         ff.push_back({f[v], v});
  39.     }
  40.  
  41.     ff.push_back({f2[p], p});
  42.     save.push_back({dp2[p], p});
  43.  
  44.     sort(save.begin(), save.end(), greater <pii> ());
  45.     sort(ff.begin(), ff.end(), greater <pii> ());
  46.     ll sum = 0;
  47.     for (int i = 0; i < min(3, (int)save.size()); i ++)
  48.         sum += save[i].st;
  49.  
  50.     for (int v: d[u]) if (v != p) {
  51.         ll cur = sum + w[u]; bool ok = false;
  52.         for (int i = 0; i < 3; i ++)
  53.             if (save.size() > i && v == save[i].nd) {
  54.                 cur -= save[i].st;
  55.                 ok = true;
  56.             }   
  57.  
  58.         if (!ok && save.size() > 2) cur -= save[2].st;
  59.  
  60.         if (ff[0].nd == v && ff.size() > 1) cur = max(cur, ff[1].st);
  61.         if (ff[0].nd != v) cur = max(cur, ff[0].st);
  62.  
  63.         f2[u] = cur;
  64.  
  65.         ans = max(ans, cur + f[v]); 
  66.  
  67.         dp2[u] = dp2[p] + w[u];
  68.         if (v == save[0].nd && save.size() > 1) dp2[u] = max(dp2[u], save[1].st + w[u]);
  69.         if (v != save[0].nd) dp2[u] = max(dp2[u], save[0].st + w[u]);
  70.         dfs2(u, v); 
  71.     }
  72.  
  73. }
  74.  
  75. int main()
  76. {
  77.     ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
  78.     // #ifndef ONLINE_JUDGE
  79.         freopen(file".inp", "r", stdin); freopen(file".out", "w", stdout);
  80.     // #endif
  81.  
  82.     cin >> n;
  83.     for (int i = 1; i <= n; i ++) cin >> w[i];
  84.     for (int i = 1, u, v; i < n; i ++) {
  85.         cin >> u >> v;
  86.         d[u].push_back(v);
  87.         d[v].push_back(u);
  88.     }
  89.  
  90.     dfs1(0, 1);
  91.     dfs2(0, 1);
  92.  
  93.     cout << ans;
  94.  
  95.     return 0;
  96. }
Success #stdin #stdout 0.01s 27316KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
Standard output is empty
https://ideone.com/cBqWDk
language:
C++14 (gcc 8.3)
created:
1 year 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!