fork download
copy 
  1. /*
  2. QUESTION 15: Google OA — Minimizing Tree Contribution by Removing Leaves
  3.  
  4. Problem Statement:
  5. You are given a rooted tree.
  6. Each node has a value.
  7.  
  8. Each remaining node contributes:
  9. value[node] * depth[node]
  10.  
  11. You may remove at most K nodes.
  12.  
  13. A node can be removed only if it is currently a leaf.
  14. After removing some leaves, their parent may become a new leaf and can be removed later.
  15.  
  16. Return the minimum possible total contribution after at most K removals.
  17.  
  18. Hard Constraints:
  19. 1 <= N <= 2000
  20. 1 <= K <= 500
  21. 1 <= value[i] <= 1e9
  22.  
  23. Example:
  24. values = [3, 8, 2, 7, 5]
  25. edges:
  26. 1 - 2
  27. 1 - 3
  28. 2 - 4
  29. 2 - 5
  30. K = 2
  31.  
  32. Expected Output:
  33. 23
  34.  
  35. Explanation:
  36. Initial contribution:
  37. 3*1 + 8*2 + 2*2 + 7*3 + 5*3 = 59
  38.  
  39. Remove leaves 4 and 5.
  40. Saved contribution:
  41. 7*3 + 5*3 = 36
  42.  
  43. Remaining contribution:
  44. 59 - 36 = 23
  45.  
  46. Brute Force:
  47. Try every possible sequence of leaf removals up to K.
  48. Compute remaining contribution.
  49.  
  50. Time Complexity:
  51. Exponential
  52.  
  53. Optimized Approach:
  54. Maximize saved contribution instead of directly minimizing remaining contribution.
  55.  
  56. dp[u][c] =
  57. maximum contribution removable from subtree u using exactly c removals.
  58.  
  59. Merge children using knapsack DP.
  60.  
  61. If the entire subtree of u is removed, it can be removed bottom-up,
  62. so that is also a valid option.
  63.  
  64. Answer:
  65. initial total contribution - best saved contribution using at most K removals.
  66.  
  67. Time Complexity:
  68. O(N * K^2)
  69.  
  70. Space Complexity:
  71. O(N * K)
  72. */
  73.  
  74. #include <bits/stdc++.h>
  75. using namespace std;
  76.  
  77. class Solution {
  78. public:
  79.     using ll = long long;
  80.  
  81.     int n, K;
  82.  
  83.     vector<vector<int>> tree;
  84.     vector<ll> value;
  85.     vector<ll> depth;
  86.  
  87.     vector<int> subtreeSize;
  88.     vector<ll> subtreeContribution;
  89.  
  90.     vector<vector<ll>> dp;
  91.  
  92.     const ll NEG = -(1LL << 60);
  93.  
  94.     void dfs(int u, int parent) {
  95.         subtreeSize[u] = 1;
  96.         subtreeContribution[u] = value[u] * depth[u];
  97.  
  98.         dp[u].assign(K + 1, NEG);
  99.         dp[u][0] = 0;
  100.  
  101.         for (int v : tree[u]) {
  102.             if (v == parent) continue;
  103.  
  104.             depth[v] = depth[u] + 1;
  105.             dfs(v, u);
  106.  
  107.             vector<ll> next(K + 1, NEG);
  108.  
  109.             for (int usedU = 0; usedU <= min(K, subtreeSize[u]); usedU++) {
  110.                 if (dp[u][usedU] == NEG) continue;
  111.  
  112.                 for (int usedV = 0;
  113.                      usedV <= min(K - usedU, subtreeSize[v]);
  114.                      usedV++) {
  115.  
  116.                     if (dp[v][usedV] == NEG) continue;
  117.  
  118.                     next[usedU + usedV] =
  119.                         max(next[usedU + usedV],
  120.                             dp[u][usedU] + dp[v][usedV]);
  121.                 }
  122.             }
  123.  
  124.             dp[u] = next;
  125.  
  126.             subtreeSize[u] += subtreeSize[v];
  127.             subtreeContribution[u] += subtreeContribution[v];
  128.         }
  129.  
  130.         if (subtreeSize[u] <= K) {
  131.             dp[u][subtreeSize[u]] =
  132.                 max(dp[u][subtreeSize[u]], subtreeContribution[u]);
  133.         }
  134.     }
  135.  
  136.     long long minimizeTreeContribution(vector<int>& values,
  137.                                        vector<vector<int>>& edges,
  138.                                        int k) {
  139.         n = values.size();
  140.         K = k;
  141.  
  142.         value.assign(n, 0);
  143.  
  144.         for (int i = 0; i < n; i++) {
  145.             value[i] = values[i];
  146.         }
  147.  
  148.         tree.assign(n, {});
  149.  
  150.         for (auto &e : edges) {
  151.             int u = e[0] - 1;
  152.             int v = e[1] - 1;
  153.  
  154.             tree[u].push_back(v);
  155.             tree[v].push_back(u);
  156.         }
  157.  
  158.         depth.assign(n, 1);
  159.         subtreeSize.assign(n, 0);
  160.         subtreeContribution.assign(n, 0);
  161.         dp.resize(n);
  162.  
  163.         dfs(0, -1);
  164.  
  165.         ll total = subtreeContribution[0];
  166.         ll bestSaved = 0;
  167.  
  168.         for (int c = 0; c <= K; c++) {
  169.             bestSaved = max(bestSaved, dp[0][c]);
  170.         }
  171.  
  172.         return total - bestSaved;
  173.     }
  174. };
  175.  
  176. int main() {
  177.     Solution sol;
  178.  
  179.     vector<int> values = {3, 8, 2, 7, 5};
  180.  
  181.     vector<vector<int>> edges = {
  182.         {1, 2},
  183.         {1, 3},
  184.         {2, 4},
  185.         {2, 5}
  186.     };
  187.  
  188.     int K = 2;
  189.  
  190.     cout << sol.minimizeTreeContribution(values, edges, K) << endl;
  191.  
  192.     return 0;
  193. }
Success #stdin #stdout 0s 5308KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
23
https://ideone.com/jibeEr
language:
C++14 (gcc 8.3)
created:
9 hours 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!