競技プログラミング日記

主に AtCoder の記事です

AtCoder Beginner Contest 259F

ABC259F

木DP. DFSは,葉から決まるので木の問題を解くのに適している. 今いる頂点 \(cu\) に対して, \(ne \in to[cu]\) の結果を集計をどうするかを考える. 今回の問題では,辺と次数が重要なので, \(cu,ne\) を結ぶ辺が影響を与える. 逆に,それ以外の辺は集計に影響を与えない.
DPは全探索であるから,素直に全探索を考えると, \(\,(cu,ne)\) の辺を使うか使わないかで 場合分けになる. まず考えると, \(ne\) から出ている辺のうち,使う本数で場合分けだろうか? しかし,実際には場合分けはまとめることができる.
辺 \(\,(cu,ne)\) を使う場合は, \(ne\) 以下の subtree において, \(ne\) から出る辺は \(d_{ne}\) 未満になる. 使わない場合は, \(d_{ne}\) 以下になる. それらの答えを \(dp_{<, ne}, dp_{\leq, ne}\) とする. 負のコストの辺があるので,\(d_{ne}\) 本全てを使わないほうが得になるケースがある. ここで,常に\(dp_{<, ne} \leq dp_{\leq, ne}\)が成立することに注意.
次に,\(ne \in to[cu]\) を動かしたときの\(\,(cu,ne)\)の辺は,どれを使えば良いのかを考える. 暫定的に全て \(dp_{\leq, ne}\) としておいて, 辺 \(\,(cu,ne)\) を選んだときに差分 \(e_{ne}\) が増えると考えればよい. つまり, \(e_{ne}\) の大きい順に辺 \(\,(cu,ne)\) を使えばよい. ただし,0以下のコストの辺は使わない方が絶対得.

実装
\(d_{v} = 0\) のとき, \(d_{v}\) 本未満は選べないので, \(dp_{<, v} = - \infty\) となる.

使っている記号,マクロ等 "https://ecsmtlir.hatenablog.com/entry/2022/12/23/131925"


int main() {
  ll n,m; cin >> n ;
  m = n-1;
  vll d(n); cinv(d);

  vector<vector<pll>> to(n);
  rep(i,m){
    ll c, a, b; cin >> a >> b >> c;
    --a, --b;
    to[a].emplace_back(b,c);
    to[b].emplace_back(a,c);
  }

  vvll dp(n, vll(2));

  auto dfs = [&](auto dfs, ll cu, ll pa = -1) -> void {
    vll v;

    for(auto [ne, cost]: to[cu]) if(ne != pa){
      dfs(dfs, ne, cu);
     
      v.push_back(- dp[ne][1] + (dp[ne][0] + cost));
      dp[cu][0] += dp[ne][1];
      dp[cu][1] += dp[ne][1];
    }
    sort(v.rbegin(), v.rend());

    rep(i,v.size()){
      if(v[i] <= 0) break;
      if(i < d[cu]-1) dp[cu][0] += v[i];
      if(i < d[cu]) dp[cu][1] += v[i];
    }
    if(d[cu] == 0) dp[cu][0] = -INFL;

    assert(dp[cu][0] <= dp[cu][1]);
    assert(dp[cu][1] >= 0);
  };
 
  dfs(dfs, 0);
  cout << dp[0][1] << endl;
 
  return 0;
}