Official

F - Select Edges Editorial by leaf1415


与えられた木を適当な頂点を根とした根付き木とみなし、動的計画法(特に、いわゆる木 DP )によってこの問題を解きます。 頂点 \(v\) 以下の部分木を単に「部分木 \(v\) 」と呼びます。

DP テーブルとして、各頂点 \(v\) について \(\mathrm{dp}_{\leq}[v]\)\(\mathrm{dp}_{\lt}[v]\) を下記の通りに定めます。

  • \(\mathrm{dp}_{\leq}[v] = \)(部分木 \(v\) のみからいくつか辺を選び、かつ、\(v\) に接続する辺は \(d_v\)以下しか選ばないときの、選ぶ辺の重みの合計としてあり得る最大値)
  • \(\mathrm{dp}_{\lt}[v] = \)(部分木 \(v\) のみからいくつか辺を選び、かつ、\(v\) に接続する辺は \(d_v\)未満しか選ばないときの、選ぶ辺の重みの合計としてあり得る最大値)

頂点 \(v\) を固定し、部分木 \(v\) のみを考えます。 \(v\) の子ら \(u_1, u_2, \ldots, u_k\)\(\mathrm{dp}_{\leq}[\ast], \mathrm{dp}_{\lt}[\ast]\) が既知のとき、\(\mathrm{dp}_{\leq}[v]\)\(\mathrm{dp}_{\lt}[v]\) を求める方法がわかれば十分です。 ここでは \(\mathrm{dp}_{\leq}[v]\) の求め方のみを述べます( \(\mathrm{dp}_{\lt}[v]\) についても「 \(d_v\) 以下」の部分を「 \(d_v\) 未満」に置き換えることで同様の議論が成立します)。

\(v\) とそれぞれの子 \(u_i\) を結ぶ辺 \((v, u_i)\) について、その辺を選ぶか選ばないかという選択肢があります。両者の場合について「部分木 \(u_i\) および辺 \((v, u_i)\) から選ぶ辺の重みの合計の最大値 \(S_i\) 」は、

  1. \((v, u_i)\) を選ぶ場合:辺 \((v, u_i)\) の重み \(w_i\) が加算された上で、\(u_i\) に接続する辺を \(d_{u_i}\) 本未満の範囲で選ぶことができるので \(S_i = \mathrm{dp}_{\lt}[u_i] + w_i\) です。
  2. \((v, u_i)\) を選ばない場合:\(u_i\) に接続する辺を \(d_{u_i}\) 本以下の範囲で選ぶことができるので、\(S_i = \mathrm{dp}_{\leq}[u_i]\) です。

よって、\(\mathrm{dp}_{\leq}[v]\) は下記の問題 \((\star)\) の答えとなります。

\((\star)\) \(v\) のそれぞれの子 \(u_i\) について、上記の 1. または 2. を選ぶ。ただし、\(1.\) を選ぶ回数は \(d_v\) 以下でなければならない。このときの選ぶ \(S_i\) の和としてあり得る最大値を求めよ。

ある子 \(u_i\) について 2. を選択しているとき、その選択を 2. から 1. に変更すると選ぶ \(S_i\) の和が \(\Delta_i := \mathrm{dp}_{\lt}[u_i] + w_i - \mathrm{dp}_{\leq}[u_i]\) だけ増加します。 よって、上記の問題 \((\star)\) を解くには、 すべての子について 2. を選んだ状態から始め「 \(v\) の子のうち \(\Delta_i\) が大きい子 \(u_i\) から貪欲に選択を 2. から 1. に変更していくこと」を、\(\Delta_i \gt 0\) かつ 1. を選ぶ子の個数が \(d_v\) 以下となる限り続けていけば良いです。

ただし上の議論の例外として、\(d_v = 0\) のときは、\(v\) に接続する辺を \(0\)未満選ぶことは出来ないので便宜上 \(\mathrm{dp}_{\lt}[v] := -\infty\) とします。

以上で本問題を \(\mathrm{O}(N \log N)\) 時間で解くことができます。 各頂点 \(v\) について、その子らを \(\Delta_i\) をキーとしてソートする操作がボトルネックとなります。

C++ 言語による正解例を以下に記載します。

#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: