Advertisement
Josif_tepe

Untitled

Dec 2nd, 2023
626
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 1.56 KB | None | 0 0
  1. #include <iostream>
  2. using namespace std;
  3. typedef long long ll;
  4. const int maxn = 5e4 + 10;
  5. const int maxk = 155;
  6. const ll INF = 1e18;
  7. int n, k;
  8. vector<int> graph[maxn];
  9. ll v[maxn];
  10. vector<int> top_sort;
  11. int subtree_size[maxn];
  12. ll dp[maxn][maxk];
  13. void dfs(int node, int parent) {
  14.     subtree_size[node] = 1;
  15.     for(int neighbour : graph[node]) {
  16.         if(neighbour != parent) {
  17.             dfs(neighbour, node);
  18.             subtree_size[node] += subtree_size[neighbour];
  19.         }
  20.     }
  21.     top_sort.push_back(node);
  22. }
  23. int main()
  24. {
  25.     ios_base::sync_with_stdio(false);
  26.     cin >> n >> k;
  27.    
  28.     for(int i = 0; i < n; i++) {
  29.         cin >> v[i];
  30.     }
  31.     for(int i = 1; i < n; i++) {
  32.         int a, b;
  33.         cin >> a >> b;
  34.         a--; b--;
  35.        
  36.         graph[a].push_back(b);
  37.         graph[b].push_back(a);
  38.     }
  39.     dfs(0, -1);
  40.     reverse(top_sort.begin(), top_sort.end());
  41.     for(int i = 0; i < maxn; i++) {
  42.         for(int j = 0; j < maxk; j++) {
  43.             dp[i][j] = -INF;
  44.         }
  45.     }
  46.     dp[0][k] = 0;
  47.    
  48.     for(int i = 0; i < n; i++) {
  49.         int node = top_sort[i];
  50.         for(int j = 0; j <= k; j++) {
  51.             if(dp[i][j] == -INF) continue;
  52.             dp[i + 1][j] = max(dp[i + 1][j], dp[i][j] + v[node]);
  53.             if(j >= 1) {
  54.                 dp[i + subtree_size[node]][j - 1] = max(dp[i + subtree_size[node]][j - 1], dp[i][j]);
  55.             }
  56.         }
  57.     }
  58.     ll res = -INF;
  59.     for(int j = 0; j <= k; j++) {
  60.         res = max(res, dp[n][j]);
  61.     }
  62.     cout << res << endl;
  63.     return 0;
  64. }
  65.  
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement