[ABC348] E - Minimize Sum of Distances

[Q] E - Minimize Sum of Distances

考察
1. とりあえず0を根とした場合の、f(0)のスコアを求める。
2. 0の子供1がいて、1に移動した場合のスコアは、なんらかの差分計算で求められそう。何の差分をとればいい?
3. 子供1に近づいた場合、子供1以下のCの総和だけスコアが下がることが分かる。また、逆に子供1に属さない頂点のCの総和だけスコアが上がってしまう。
4. ということは、
子供1以下のCの総和 > 全体Cの総和 / 2
なら、子供1に移動したほうがスコアが下がって嬉しい。
改善がなくなるまで貪欲的に繰り返せばよさそう。

入力例1で整理する

1. 0を根として、スコアとコストを出す。
コスト = 各頂点vについて v以下のCの総和
スコア = 距離 * コスト

求めるとこんな感じ。
score[0]:6 costs[0]:5
score[1]:5 costs[1]:3
score[2]:1 costs[2]:1
score[3]:2 costs[3]:2

2. G[0]の子供に遷移した場合、スコアの改善があるかシミュレートする。

total_cost = costs[0] = 5がCの総和。

G[0] = {1, 2}なので、子供1と、子供2に遷移した場合のスコアの変化を見る。
score[0]:6 costs[0]:5
score[1]:5 costs[1]:3 ... 0->1に行くと、scoreが costs[1]=3 下がって、5-3=2 上がる。結果2-3=-1でスコアが下がるから嬉しい。
score[2]:1 costs[2]:1 ... 0->2に行くと、scoreが costa[2]=1 下がって、5-1=4 上がる。結果4-1=3でスコアが上がるからダメ
score[3]:2 costs[3]:2

頂点1に移動してスコアが-1され、5が得られる。
さらに、G[1] = {0, 3}なので、孫3 に移動した場合のスコア変化を見る。
score[0]:6 costs[0]:5 ...0には戻らない。
score[1]:5 costs[1]:3 
score[2]:1 costs[2]:1 
score[3]:2 costs[3]:2 ... 1->3に行くと、scoreが costs[3]=2 下がって、5-2=3 上がる。結果3-2=1でスコアが上がるのでダメ。

3に移動するとスコアが上がってしまうためおしまい。
f(1) = 5が答え

実装

int main()
{
  cincout();

  ll N;
  cin >> N;
  vector<vector<ll>> G(N);
  vector<ll> C(N);
  rep(i, N - 1)
  {
    ll a, b;
    cin >> a >> b;
    --a, --b;
    G[a].push_back(b);
    G[b].push_back(a);
  }
  rep(i, N) cin >> C[i];

  vector<ll> scores(N), costs(N); // 0を根としたとき score:f(v) cost:Cの総和
  function<pll(ll, ll)> init = [&](ll v, ll p) -> pll
  {
    ll cost = 0;
    ll score = 0;
    for (auto u : G[v])
    {
      if (u == p)
        continue;
      auto [c, s] = init(u, v);
      cost += c;
      score += s;
    }
    cost += C[v];
    if (v != 0){ // 0を根としたときのスコアが欲しいので、0の場合はスコアを足さない
      score += cost;
    }
    scores[v] = score;
    costs[v] = cost;
    return {cost, score};
  };
  init(0, -1);
  ll score = scores[0];
  ll cur = 0;
  ll total_cost = costs[0];
  ll p = -1;
  while (1)
  {
    bool is_update = false;
    for (auto u : G[cur])
    {
      if (u == p) continue;
      if (costs[u] > total_cost / 2) // 近づいたらスコアが下がるなら行く
      {
        score = score - costs[u] + (total_cost - costs[u]);
        p = cur;
        cur = u;
        is_update = true;
        break;
      }
    }
    if (!is_update)
      break;
  }
  cout << score << endl;
}

https://atcoder.jp/contests/abc348/submissions/52118085

この記事が気に入ったらサポートをしてみませんか?