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\) 」は、
- 辺 \((v, u_i)\) を選ぶ場合:辺 \((v, u_i)\) の重み \(w_i\) が加算された上で、\(u_i\) に接続する辺を \(d_{u_i}\) 本未満の範囲で選ぶことができるので \(S_i = \mathrm{dp}_{\lt}[u_i] + w_i\) です。
- 辺 \((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: