木を頂点 \(1\) を根とした根付き木であるとします。
\(seq\) を頂点 \(1\) を始点として BFS した際の頂点順とします。

最初の BFS の際に、その頂点の親および親に向かう辺の番号を求めておくことができます。

全ての辺について、その辺を取り除いた場合の \(2\) つの連結成分の大きさは、 \(seq\) を後ろから前に調べることで、各頂点の部分木の大きさとして葉から順に求められます。

ここから、実際に本問題で要求されている加算を行います。ここで、加算の順序を入れ替えてもよい(辺を任意の順序で処理してよい)ことに注意してください。

頂点 \(v\) から根に向かって辺 \(i\) が伸びているとします。
もし \(v\) 以下の部分木の大きさが \(n/2\) 以上である場合、 \(v\) 以下の部分木全体に \(i\) を加算すればよいです。
\(v\) 以下の部分木の大きさが \(n/2\) 以下である場合も、「部分木全体への加算」として扱うことができます。その方法とは、全ての頂点に \(i\) を加算すると取り決めたうえで、 \(v\) 以下の部分木全体に \(-i\) を加算するということです。

部分木全体への加算は、その頂点に値を加算したうえで \(seq\) を前から調べることで葉に向かって加算を伝搬させれば実現可能です。

一連の解法は時間計算量 \(O(N)\) で動作します。

実装例 (C++):

#include<bits/stdc++.h>

using namespace std;
using ll=long long;
using pi=pair<int,int>;
using Graph=vector<vector<pi>>;

int main(){
  ios::sync_with_stdio(false);
  cin.tie(nullptr);
  int n;
  cin >> n;
  Graph g(n);
  for(int i=1;i<n;i++){
    int a,b;
    cin >> a >> b;
    a--; b--;
    g[a].push_back({b,i});
    g[b].push_back({a,i});
  }
  vector<int> seq;
  vector<int> par(n,-1);
  vector<int> peg(n);
  {
    queue<int> q;
    par[0]=-2;
    peg[0]=0;
    q.push(0);
    while(!q.empty()){
      auto od=q.front(); q.pop();
      seq.push_back(od);
      for(auto &nx : g[od]){
        if(par[nx.first]==-1){
          par[nx.first]=od;
          peg[nx.first]=nx.second;
          q.push(nx.first);
        }
      }
    }
  }

  reverse(seq.begin(),seq.end());
  vector<int> sz(n,1);
  vector<ll> x(n,0);
  ll glob=0;
  for(auto &nx : seq){
    for(auto &ny : g[nx]){
      if(ny.first==par[nx]){continue;}
      sz[nx]+=sz[ny.first];
    }
    if(2*sz[nx]<n){
      glob+=peg[nx];
      x[nx]-=peg[nx];
    }
    else{
      x[nx]+=peg[nx];
    }
  }

  reverse(seq.begin(),seq.end());
  for(auto &nx : seq){
    for(auto &ny : g[nx]){
      if(ny.first==par[nx]){continue;}
      x[ny.first]+=x[nx];
    }
  }

  for(int i=0;i<n;i++){
    if(i){cout << " ";}
    cout << glob+x[i];
  }cout << "\n";
  return 0;
}

原案: admin