Advertisement
rishu110067

Untitled

Jan 22nd, 2025
59
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 1.69 KB | None | 0 0
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3.  
  4. const int N = 101;
  5. int K;
  6. vector<int> tree(N);
  7. vector<pair<int, int>> v[N];
  8. int dp[N][N];
  9. int countTree[N][N];
  10.  
  11. int solve(int x, int k) {
  12.     if (dp[x][k] != -1) {
  13.         return dp[x][k];
  14.     }
  15.    
  16.     // don't make this node sawmill
  17.     int cost1 = 0;
  18.     int tree1 = tree[x];
  19.     for (auto [i, d]: v[x]) {
  20.         cost1 += solve(i, k);
  21.         cost1 += d * countTree[i][k];
  22.         tree1 += countTree[i][k];
  23.     }
  24.    
  25.     if (k == K || x == 0) {
  26.         countTree[x][k] = tree1;
  27.         cout << x << " " << k << " -> [10] " << cost1 << " " << tree1 << endl;
  28.         return dp[x][k] = cost1;
  29.     }
  30.    
  31.     // make this node sawmill
  32.     int cost2 = 0;
  33.     int tree2 = 0;
  34.     for (auto [i, d]: v[x]) {
  35.         cost2 += solve(i, k + 1);
  36.         cost2 += d * countTree[i][k + 1];
  37.         tree2 += countTree[i][k + 1];
  38.     }
  39.    
  40.     if (cost1 < cost2 || (cost1 == cost2 && tree1 <= tree2)) {
  41.         countTree[x][k] = tree1;
  42.         cout << x << " " << k << " -> [1] "  << cost1 << " " << tree1 << endl;
  43.         return dp[x][k] = cost1;
  44.     } else {
  45.         countTree[x][k] = tree2;
  46.         cout << x << " " << k << " -> [2] " << cost2 << " " << tree2 << endl;
  47.         return dp[x][k] = cost2;
  48.     }
  49. }
  50.  
  51. int main() {
  52.     int n;
  53.     cin >> n >> K;
  54.     for (int i = 0; i < n; i++) {
  55.         cin >> tree[i];
  56.     }
  57.     for (int i = 0; i < n-1; i++) {
  58.         int x, y, w;
  59.         cin >> x >> y >> w;
  60.         v[y].push_back({x, w});
  61.     }
  62.     for (int i = 0; i < n; i++) {
  63.         for (int j = 0; j <= K; j++) {
  64.             dp[i][j] = -1;
  65.             countTree[i][j] = -1;
  66.         }
  67.     }
  68.     cout << solve(0, 0) << endl;
  69. }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement