[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;
}