Official

F - Select Edges Editorial by en_translator


We regard the given tree as a rooted tree rooted at some arbitrary vertex and solve this problem with Dynamic Programming (specifically, so-called Tree DP). We call the subtree rooted at Vertex \(v\) simply “Subtree \(v\).”

For each vertex \(v\), defined the DP table \(\mathrm{dp}_{\leq}[v]\) and \(\mathrm{dp}_{\lt}[v]\) as follows.

  • \(\mathrm{dp}_{\leq}[v] = \) (the maximum sum of weights when choosing some of the edges from Subtree \(v\) so that at most \(d_v\) edges incident to \(v\) is chosen)
  • \(\mathrm{dp}_{\lt}[v] = \) (the maximum sum of weights when choosing some of the edges from Subtree \(v\) so that strictly less than \(d_v\) edges incident to \(v\) is chosen)

We now consider Subtree \(v\) for some fixed vertex \(v\). All we want is how to find \(\mathrm{dp}_{\leq}[v]\) and \(\mathrm{dp}_{\lt}[v]\) when \(\mathrm{dp}_{\leq}[\ast], \mathrm{dp}_{\lt}[\ast]\) for the children \(u_1, u_2, \ldots, u_k\) of \(v\) are known. Here, we only describe how to find \(\mathrm{dp}_{\leq}[v]\). (The \(\mathrm{dp}_{\lt}[v]\) part can be discussed similarly by replacing “at most” with “strictly less than \(d_v\).”

For each edge \((v, u_i)\) connecting \(v\) and its \(u_i\), we have two choice: to choose, or not to choose, the edge. For each case, “the maximum sum of weights \(S_i\) chosen from Subtree \(u_i\) and edge \((v, u_i)\)” is:

  1. if edge \((v, u_i)\) is chosen: the weight \(w_i\) of the edge \((v, u_i)\), and less than \(d_{u_i}\) more edges can be chosen incident to \(u_i\), so the answer is \(S_i = \mathrm{dp}_{\lt}[u_i] + w_i\).
  2. if edge \((v, u_i)\) is not chosen: at most \(d_{u_i}\) more edges incident to \(u_i\) can be chosen, so \(S_i = \mathrm{dp}_{\leq}[u_i]\).

Thus, the answer is equal to the answer to the following problem \((\star)\).

\((\star)\): for each children \(u_i\) of \(v\), you choose either 1. or 2. above. Here, 1. cam be chosen at most \(d_v\) times. Find the maximum possible sum of \(S_i\) chosen.

If 2. is chosen for a children, changing the choice from 2. to 1. changes the sum of \(S_i\) by \(\Delta_i := \mathrm{dp}_{\lt}[u_i] + w_i - \mathrm{dp}_{\leq}[u_i]\). Thus, in order to solve the problem \((\star)\) above, you start from the state where 2. is chosen for all children, and “change the choice from 2. to 1. greedily, with the children \(u_i\) with larger \(\Delta_i\) first” while \(\Delta_i \gt 0\) and 1. is chosen at most \(d_v\) times.

The exception of the discussion above is when \(d_v = 0\), where you cannot choose strictly less than 0 edges incident to \(v\), so for simplicity, we let \(\mathrm{dp}_{\lt} := -\infty\).

Therefore, the problem has been solved in a total of \(\mathrm{O}(N \log N)\) time. The bottleneck is sorting the children of each vertex \(v\) by the keys \(\Delta_i\).

The following is a sample code in C++ language.

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;
typedef long long ll;
typedef pair<int, ll> E;
const ll inf = 1e18;

int n;
int d[300005];
vector<E> G[300005];
ll dp[300005][2];

void dfs(int v, int p)
{
  vector<ll> vec;
  for(auto e : G[v]){
    int u = e.first, w = e.second;
    if(u == p) continue;
    dfs(u, v);
    vec.push_back(dp[u][0]+w-dp[u][1]);
    dp[v][0] += dp[u][1], dp[v][1] += dp[u][1];
  }
  sort(vec.rbegin(), vec.rend());
  
  for(int i = 0; i < (int)vec.size(); i++){
    if(vec[i] <= 0) break;
    if(i < d[v]-1) dp[v][0] += vec[i];
    if(i < d[v]) dp[v][1] += vec[i];
  }
  if(d[v] == 0) dp[v][0] = -inf;
}

int main(void)
{
  cin >> n;
  for(int i = 1; i <= n; i++) cin >> d[i];
  int u, v, w;
  for(int i = 1; i <= n-1; i++){
    cin >> u >> v >> w;
    G[u].push_back(E(v, w));
    G[v].push_back(E(u, w));
  }
  dfs(1, -1);
  cout << dp[1][1] << endl;
    
  return 0;
}

posted:
last update: