API tools faq
paste
Advertisement
ThegeekKnight16

Chase

ThegeekKnight16
Oct 14th, 2023
731
0
Never
Add comment
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 3.82 KB | None | 0 0
raw download clone embed print report
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3. #pragma GCC optimize("Ofast,unroll-loops")
  4. #pragma GCC target("avx2,tune=native")
  5. #define int long long
  6. const int MAXN = 1e5 + 10;
  7. const int MAXV = 101;
  8. array<array<array<int, 3>, MAXV>, MAXN> dp;
  9. array<vector<int>, MAXN> grafo;
  10. array<int, MAXN> pom, pomviz;
  11. int N, V;
  12. int resp = 0;
  13.  
  14. void dfs(int v, int p)
  15. {
  16.     dp[v].fill({0, 0, 0});
  17.     pomviz[v] = 0;
  18.     for (auto viz : grafo[v])
  19.     {
  20.         if (viz == p) continue;
  21.         pomviz[v] += pom[viz];
  22.         dfs(viz, v);
  23.     }
  24.  
  25.     for (auto viz : grafo[v])
  26.     {
  27.         if (viz == p) continue;
  28.         for (int k = 1; k <= V; k++)
  29.         {
  30.             int prox = max(dp[viz][k][0], dp[viz][k-1][0] + pomviz[v]);
  31.             auto &atual = dp[v][k];
  32.             if (prox > atual[0])
  33.             {
  34.                 swap(atual[2], atual[1]);
  35.                 swap(atual[1], atual[0]);
  36.                 atual[0] = prox;
  37.             }
  38.             else if (prox > atual[1])
  39.             {
  40.                 swap(atual[2], atual[1]);
  41.                 atual[1] = prox;
  42.             }
  43.             else if (prox > atual[2])
  44.             {
  45.                 atual[2] = prox;
  46.             }
  47.         }
  48.     }
  49. }
  50.  
  51. void girarArvore(int v, int p)
  52. {
  53.     // cerr << "v: " << v << '\n';
  54.     // cerr << '\t' << "dp: ";
  55.     // for (auto x : dp[v][V]) cerr << get<0>(x) << " ";
  56.     // cerr << '\n';
  57.     // cerr << '\t' << "pomviz: " << pomviz[v] << '\n';
  58.     resp = max(resp, dp[v][V][0]);
  59.  
  60.     for (auto viz : grafo[v])
  61.     {
  62.         if (viz == p) continue;
  63.  
  64.         //Salva
  65.         auto antv = dp[v]; int antpom = pomviz[v];
  66.         auto antviz = dp[viz]; int antpomviz = pomviz[viz];
  67.  
  68.         //Gira
  69.         pomviz[v] -= pom[viz]; pomviz[viz] += pom[v];
  70.         dp[v].fill({0, 0, 0}); dp[viz].fill({0, 0, 0});
  71.         for (auto viz2 : grafo[v])
  72.         {
  73.             for (int k = 1; k <= V; k++)
  74.             {
  75.                 auto &atual = dp[v][k];
  76.                 if (viz2 == viz) continue;
  77.                 int prox = max(dp[viz2][k][0], dp[viz2][k-1][0] + pomviz[v]);
  78.                 if (prox > atual[0])
  79.                 {
  80.                     swap(atual[2], atual[1]);
  81.                     swap(atual[1], atual[0]);
  82.                     atual[0] = prox;
  83.                 }
  84.                 else if (prox > atual[1])
  85.                 {
  86.                     swap(atual[2], atual[1]);
  87.                     atual[1] = prox;
  88.                 }
  89.                 else if (prox > atual[2])
  90.                 {
  91.                     atual[2] = prox;
  92.                 }
  93.             }
  94.         }
  95.         for (auto viz2 : grafo[viz])
  96.         {
  97.             for (int k = 1; k <= V; k++)
  98.             {
  99.                 auto &atual = dp[viz][k];
  100.                 int prox = max(dp[viz2][k][0], dp[viz2][k-1][0] + pomviz[viz]);
  101.                 if (prox > atual[0])
  102.                 {
  103.                     swap(atual[2], atual[1]);
  104.                     swap(atual[1], atual[0]);
  105.                     atual[0] = prox;
  106.                 }
  107.                 else if (prox > atual[1])
  108.                 {
  109.                     swap(atual[2], atual[1]);
  110.                     atual[1] = prox;
  111.                 }
  112.                 else if (prox > atual[2])
  113.                 {
  114.                     atual[2] = prox;
  115.                 }
  116.             }
  117.         }
  118.  
  119.         girarArvore(viz, v);
  120.  
  121.         //Volta
  122.         dp[v] = antv; dp[viz] = antviz;
  123.         pomviz[v] = antpom; pomviz[viz] = antpomviz;
  124.     }
  125. }
  126.  
  127. int32_t main()
  128. {
  129.     ios_base::sync_with_stdio(false);
  130.     cin.tie(NULL);
  131.     cin >> N >> V;
  132.     for (int i = 1; i <= N; i++) cin >> pom[i];
  133.     for (int i = 1; i < N; i++)
  134.     {
  135.         int X, Y;
  136.         cin >> X >> Y;
  137.         grafo[X].push_back(Y);
  138.         grafo[Y].push_back(X);
  139.     }
  140.     dfs(1, 0);
  141.     girarArvore(1, 1);
  142.     cout << resp << '\n';
  143. }
  144.  
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement
Public Pastes
Advertisement
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy.  OK, I Understand
Not a member of Pastebin yet?
Sign Up, it unlocks many cool features!