fork download
copy 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3.  
  4. /*
  5.     Problem: Maximum Weight Forest with Component Size Limit
  6.  
  7.     You are given an undirected weighted tree with n nodes numbered from 0 to n - 1.
  8.  
  9.     Each edge is given as:
  10.         [u, v, w]
  11.     meaning there is an edge between u and v with weight w.
  12.  
  13.     You are allowed to remove any number of edges.
  14.     After removing edges, the tree becomes a forest.
  15.  
  16.     A forest is considered valid if every connected component
  17.     has size at most k.
  18.  
  19.     Your task is to return the maximum possible sum of weights
  20.     of the edges that remain, such that the resulting forest is valid.
  21.  
  22.     ------------------------------------------------------------
  23.  
  24.     Example 1:
  25.     n = 3, k = 2
  26.     edges = {{0,1,5}, {1,2,10}}
  27.  
  28.     If we keep both edges, all 3 nodes stay connected,
  29.     so component size becomes 3, which is not allowed.
  30.  
  31.     We must remove one edge.
  32.     Best choice is to keep edge (1,2) with weight 10.
  33.  
  34.     Answer = 10
  35.  
  36.     ------------------------------------------------------------
  37.  
  38.     Example 2:
  39.     n = 5, k = 2
  40.     edges = {{0,1,4}, {1,2,3}, {1,3,7}, {3,4,2}}
  41.  
  42.     We need every connected component to contain at most 2 nodes.
  43.     So we remove some edges in a way that keeps as much total weight as possible.
  44.  
  45.     ------------------------------------------------------------
  46.  
  47.     Main idea:
  48.  
  49.     This is a tree DP problem.
  50.  
  51.     Let dp[u][s] mean:
  52.  
  53.     - we are only considering the subtree rooted at u
  54.     - all connected components inside this subtree are already valid
  55.     - the component that contains u currently has exactly s nodes
  56.     - dp[u][s] = maximum total weight we can keep inside this subtree
  57.  
  58.     Now process children one by one.
  59.  
  60.     For each child v of u, we have 2 choices:
  61.  
  62.     1) Cut the edge (u, v)
  63.        Then v's subtree becomes a separate component.
  64.        So we just add the best valid answer from v's subtree.
  65.  
  66.     2) Keep the edge (u, v)
  67.        Then the component containing u merges with the component containing v.
  68.        If u-part has size s and v-part has size t,
  69.        then merged size becomes s + t.
  70.        This is allowed only if s + t <= k.
  71.  
  72.     So this becomes a knapsack-style merge on tree DP.
  73.  
  74.     Time Complexity:
  75.     Roughly O(n * k^2)
  76.  
  77.     Space Complexity:
  78.     O(n * k)
  79. */
  80.  
  81. class Solution {
  82. public:
  83.     long long maximumKeptWeight(int n, int k, vector<vector<int>>& edges) {
  84.         vector<vector<pair<int, int>>> tree(n);
  85.  
  86.         for (auto& e : edges) {
  87.             int u = e[0];
  88.             int v = e[1];
  89.             int w = e[2];
  90.             tree[u].push_back({v, w});
  91.             tree[v].push_back({u, w});
  92.         }
  93.  
  94.         const long long NEG = -(1LL << 60);
  95.  
  96.         // dp[u][s] = best answer for subtree of u,
  97.         // when the component containing u has size s
  98.         vector<vector<long long>> dp(n);
  99.  
  100.         // subtreeSize[u] = number of nodes in subtree of u
  101.         vector<int> subtreeSize(n, 0);
  102.  
  103.         function<void(int, int)> dfs = [&](int u, int parent) {
  104.             subtreeSize[u] = 1;
  105.  
  106.             // Initially only node u is present
  107.             dp[u] = vector<long long>(2, NEG);
  108.             dp[u][1] = 0;
  109.  
  110.             for (auto& [v, w] : tree[u]) {
  111.                 if (v == parent) continue;
  112.  
  113.                 dfs(v, u);
  114.  
  115.                 int limitU = min(subtreeSize[u], k);
  116.                 int limitV = min(subtreeSize[v], k);
  117.  
  118.                 // If we cut edge (u, v), then v becomes separate.
  119.                 // We only need the best valid answer from v's subtree.
  120.                 long long bestSeparate = NEG;
  121.                 for (int t = 1; t <= limitV; t++) {
  122.                     bestSeparate = max(bestSeparate, dp[v][t]);
  123.                 }
  124.  
  125.                 // New DP after processing child v
  126.                 vector<long long> nextDp(min(subtreeSize[u] + subtreeSize[v], k) + 1, NEG);
  127.  
  128.                 for (int s = 1; s <= limitU; s++) {
  129.                     if (dp[u][s] == NEG) continue;
  130.  
  131.                     // Option 1: cut edge (u, v)
  132.                     // u's component size stays the same
  133.                     nextDp[s] = max(nextDp[s], dp[u][s] + bestSeparate);
  134.  
  135.                     // Option 2: keep edge (u, v)
  136.                     // merge u's component with v's component
  137.                     for (int t = 1; t <= limitV && s + t <= k; t++) {
  138.                         if (dp[v][t] == NEG) continue;
  139.  
  140.                         nextDp[s + t] = max(
  141.                             nextDp[s + t],
  142.                             dp[u][s] + dp[v][t] + w
  143.                         );
  144.                     }
  145.                 }
  146.  
  147.                 subtreeSize[u] += subtreeSize[v];
  148.                 dp[u].swap(nextDp);
  149.             }
  150.         };
  151.  
  152.         // Root the tree at node 0
  153.         dfs(0, -1);
  154.  
  155.         // At the root, any valid component size is fine
  156.         long long answer = 0;
  157.         for (int s = 1; s < (int)dp[0].size(); s++) {
  158.             answer = max(answer, dp[0][s]);
  159.         }
  160.  
  161.         return answer;
  162.     }
  163. };
Compilation error #stdin compilation error #stdout 0s 0KB
comments (?)
stdin
copy 
Standard input is empty
compilation info 
prog.cpp: In lambda function:
prog.cpp:110:24: warning: structured bindings only available with -std=c++17 or -std=gnu++17
             for (auto& [v, w] : tree[u]) {
                        ^
/usr/bin/ld: /usr/lib/gcc/x86_64-linux-gnu/8/../../../x86_64-linux-gnu/Scrt1.o: in function `_start':
(.text+0x20): undefined reference to `main'
collect2: error: ld returned 1 exit status
stdout
copy 
Standard output is empty
https://ideone.com/SDDpKw
language:
C++14 (gcc 8.3)
created:
4 days 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!