Advertisement
libobil

Untitled

Sep 22nd, 2023
566
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 6.80 KB | None | 0 0
  1. #include "closing.h"
  2. #include <algorithm>
  3. #include <iostream>
  4. #include <numeric>
  5. #include <cassert>
  6. #include <vector>
  7. #include <queue>
  8.  
  9. typedef long long llong;
  10. const int MAXN = 500 + 10;
  11. const llong INF = 8e18;
  12. const int INTINF = 2e9;
  13.  
  14. int n;
  15. int x;
  16. int y;
  17. llong k;
  18. int sz[MAXN];
  19. int par[MAXN];
  20. bool isOnPath[MAXN];
  21. int suffixSizes[MAXN];
  22. std::vector <int> path;
  23. std::vector <int> pathSizes;
  24. std::vector <std::pair <int,int>> g[MAXN];
  25. llong dpR[MAXN][2 * MAXN][2][2]; // root
  26. llong dpC[MAXN][2 * MAXN][2][2]; // children, knapsack
  27. llong dpS[MAXN][2 * MAXN][2][2]; // stick
  28. bool blR[MAXN][2 * MAXN][2][2];
  29. bool blC[MAXN][2 * MAXN][2][2];
  30. bool blS[MAXN][2 * MAXN][2][2];
  31. llong distX[MAXN];
  32. llong distY[MAXN];
  33.  
  34. void reset()
  35. {
  36.     path.clear();
  37.     pathSizes.clear();
  38.     for (int i = 0 ; i < n ; ++i)
  39.     {
  40.         g[i].clear();
  41.         isOnPath[i] = false;
  42.         distX[i] = distY[i] = 0;
  43.         for (int j = 0 ; j <= 2 * n ; ++j)
  44.         {
  45.             for (int flagX = 0 ; flagX < 2 ; ++flagX)
  46.             {
  47.                 for (int flagY = 0 ; flagY < 2 ; ++flagY)
  48.                 {
  49.                     dpR[i][j][flagX][flagY] = 0;
  50.                     dpC[i][j][flagX][flagY] = 0;
  51.                     dpS[i][j][flagX][flagY] = 0;
  52.                     blR[i][j][flagX][flagY] = false;
  53.                     blC[i][j][flagX][flagY] = false;
  54.                     blS[i][j][flagX][flagY] = false;
  55.                 }
  56.             }
  57.         }
  58.     }
  59. }
  60.  
  61. void dfs(int node, int par, llong to[], llong dist)
  62. {
  63.     to[node] = dist;
  64.     for (const auto &[u, d] : g[node])
  65.     {
  66.         if (u == par)
  67.         {
  68.             continue;
  69.         }
  70.  
  71.         dfs(u, node, to, dist + d);
  72.     }
  73. }
  74.  
  75. bool dfsPath(int node, int par)
  76. {
  77.     if (node == y)
  78.     {
  79.         path.push_back(node);
  80.         return true;
  81.     }
  82.  
  83.     for (const auto &[u, d] : g[node])
  84.     {
  85.         if (u == par)
  86.         {
  87.             continue;
  88.         }
  89.  
  90.         if (dfsPath(u, node))
  91.         {
  92.             path.push_back(node);
  93.             return true;
  94.         }
  95.     }
  96.  
  97.     return false;
  98. }
  99.  
  100. void dfsSize(int node, int p)
  101. {
  102.     sz[node] = 1;
  103.     par[node] = p;
  104.     for (const auto &[u, d] : g[node])
  105.     {
  106.         if (p == u)
  107.         {
  108.             continue;
  109.         }
  110.  
  111.         dfsSize(u, node);
  112.         if (!isOnPath[u]) sz[node] += sz[u];
  113.     }
  114. }
  115.  
  116. llong fRoot(int, int, int, int);
  117. llong fChildren(int, int, int, int, int);
  118. llong fStick(int, int, int, int);
  119.  
  120. llong fRoot(int node, int choose, int flagX, int flagY)
  121. {
  122.     if (choose == 0)
  123.     {
  124.         return 0;
  125.     }
  126.  
  127.     if (!flagX && !flagY)
  128.     {
  129.         return k + 1;
  130.     }
  131.  
  132.     if (blR[node][choose][flagX][flagY])
  133.     {
  134.         return dpR[node][choose][flagX][flagY];
  135.     }
  136.  
  137.     blR[node][choose][flagX][flagY] = true;
  138.     dpR[node][choose][flagX][flagY] = fChildren(node, 0, choose, flagX, flagY);
  139.     dpR[node][choose][flagX][flagY] = std::min(dpR[node][choose][flagX][flagY], k + 1);
  140.     return dpR[node][choose][flagX][flagY];
  141. }
  142.  
  143. llong fChildren(int node, int idx, int choose, int flagX, int flagY)
  144. {
  145.     if (idx == g[node].size())
  146.     {
  147.         if (choose == 0) return 0;
  148.         return k + 1;
  149.     }
  150.  
  151.     int current = g[node][idx].first;
  152.     if (isOnPath[current] || current == par[node])
  153.     {
  154.         return fChildren(node, idx + 1, choose, flagX, flagY);
  155.     }
  156.  
  157.     if (blC[current][choose][flagX][flagY])
  158.     {
  159.         return dpC[current][choose][flagX][flagY];
  160.     }
  161.    
  162.     blC[current][choose][flagX][flagY] = true;
  163.     dpC[current][choose][flagX][flagY] = k + 1;
  164.    
  165.     for (int currX = 0 ; currX <= flagX ; ++currX)
  166.     {
  167.         for (int currY = 0 ; currY <= flagY ; ++currY)
  168.         {
  169.             for (int subtreeChoose = std::max(choose - 2 * suffixSizes[(idx + 1 == g[node].size() ? 0 : g[node][idx + 1].first)], currX + currY) ; subtreeChoose <= std::min(2 * sz[g[node][idx].first], choose) ; ++subtreeChoose)
  170.             {
  171.                 if (fChildren(node, idx + 1, choose - subtreeChoose, flagX, flagY) > k) continue;
  172.                 dpC[current][choose][flagX][flagY] = std::min(dpC[current][choose][flagX][flagY], fChildren(node, idx + 1, choose - subtreeChoose, flagX, flagY) + std::max((currX == 0 ? 0 : distX[current]), (currY == 0 ? 0 : distY[current])) + fRoot(current, subtreeChoose - currX - currY, currX, currY));
  173.             }
  174.         }
  175.     }
  176.  
  177.     return dpC[current][choose][flagX][flagY];
  178. }
  179.  
  180. llong fStick(int idx, int choose, int flagX, int flagY)
  181. {
  182.     if (idx == path.size())
  183.     {
  184.         if (choose == 0) return 0;
  185.         return k + 1;
  186.     }
  187.  
  188.     if (blS[idx][choose][flagX][flagY])
  189.     {
  190.         return dpS[idx][choose][flagX][flagY];
  191.     }
  192.  
  193.     blS[idx][choose][flagX][flagY] = true;
  194.     dpS[idx][choose][flagX][flagY] = k + 1;
  195.     for (int currX = 0 ; currX <= flagX ; ++currX)
  196.     {
  197.         for (int currY = flagY ; currY < 2 ; ++currY)
  198.         {
  199.             for (int subtreeChoose = std::max(choose - 2 * pathSizes[idx + 1], currX + currY) ; subtreeChoose <= std::min(2 * sz[path[idx]], choose) ; ++subtreeChoose)
  200.             {
  201.                 if (fStick(idx + 1, choose - subtreeChoose, currX, currY) > k) continue;
  202.                 dpS[idx][choose][flagX][flagY] = std::min(dpS[idx][choose][flagX][flagY], fStick(idx + 1, choose - subtreeChoose, currX, currY) + std::max((currX == 0 ? 0 : distX[path[idx]]), (currY == 0 ? 0 : distY[path[idx]])) + fRoot(path[idx], subtreeChoose - currX - currY, currX, currY));
  203.             }
  204.         }
  205.     }
  206.  
  207.     return dpS[idx][choose][flagX][flagY];
  208. }
  209.  
  210. int max_score(int N, int X, int Y, long long K, std::vector <int> U, std::vector <int> V, std::vector <int> W)
  211. {
  212.     reset();
  213.     n = N;
  214.     x = X;
  215.     y = Y;
  216.     k = K;
  217.  
  218.     for (int i = 0 ; i < n - 1 ; ++i)
  219.     {
  220.         g[U[i]].emplace_back(V[i], W[i]);
  221.         g[V[i]].emplace_back(U[i], W[i]);
  222.     }    
  223.  
  224.     dfs(x, -1, distX, 0);
  225.     dfs(y, -1, distY, 0);
  226.     dfsPath(x, -1);
  227.     dfsSize(x, -1);
  228.  
  229.     for (int i = 0 ; i < n ; ++i)
  230.     {
  231.         for (int j = g[i].size() - 1 ; j >= 0 ; --j)
  232.         {
  233.             int add = 0;
  234.             suffixSizes[g[i][j].first] = suffixSizes[(j + 1 == g[i].size() ? 0 : g[i][j + 1].first)];
  235.             if (!isOnPath[g[i][j].first] && g[i][j].first != par[i])
  236.             {
  237.                 add = sz[g[i][j].first];
  238.             }
  239.  
  240.             suffixSizes[g[i][j].first] += add;
  241.         }
  242.     }
  243.  
  244.     std::reverse(path.begin(), path.end());
  245.     for (const int &i : path)
  246.     {
  247.         isOnPath[i] = true;
  248.     }
  249.  
  250.     pathSizes.resize(path.size() + 1, 0);
  251.     for (int i = path.size() - 1 ; i >= 0 ; --i)
  252.     {
  253.         pathSizes[i] = pathSizes[i + 1] + sz[path[i]];
  254.     }
  255.    
  256.     for (int i = 2 * n ; i >= 1 ; --i)
  257.     {
  258.         if (fStick(0, i, 1, 0) <= k)
  259.         {
  260.             return i;
  261.         }
  262.     }
  263.  
  264.     return 0;
  265. }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement