競技プログラミング日記

主に AtCoder の記事です

AtCoder Beginner Contest 220F

ABC220F

差分の更新で高速化の問題.

まず,vertex を一つ固定した場合の答えを求まる. Vertex を 0で考えると,0 を根とした木に対してDFSをすればよい. 根からの距離はDFSの深さで求まる.

次に差分を考える. \(s \rightarrow t\) と vertex を遷移したとする. このとき,距離を求める際に,大部分の計算が使いまわせると気づく.
\(V_{x}\) により,vertex \(x\) を根とした subtree の vertex の個数を表す. \(t\) に遷移したことで,\(t\) 以下のsubtree の頂点に対しては距離が -1 される. それ以外の頂点に関しては,距離が +1 される.
よって, \begin{align} ans_{t} &= ans_{s} - V_{t} + (N - V_{t}) \\ &= ans_{s} + N -2 V_{t} \end{align} となる.

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

int main() {
  ll n,m; cin >> n;
  m = n-1;
  vvll to(n);
  rep(i,m){
    ll a, b; cin >> a >> b;
    --a, --b;
    to[a].push_back(b);
    to[b].push_back(a);
  }

  vll dep(n);
  vll ans(n);
  vll t(n,1); // subtree の vx の個数
  auto dfs = [&](auto f, ll type, ll cu = 0, ll pa = -1, ll d = 0) -> void{
    dep[cu] = d;

    for(auto ne: to[cu]) if(ne != pa){
      // rem: 遷移前に済ませる
      // ans[cu] が求まった状態で ans[ne] を計算したいから.
      if(type == 1) ans[ne] = ans[cu] + n - 2*t[ne];

      f(f, type, ne, cu, d+1); // 遷移する時点で,ans[ne] が求まっている必要がある.
     
      // rem: 遷移後に済ませる
      if(type == 0) t[cu] += t[ne]; // 遷移が終わってから subtree の vx を集計
    }
   
  };

  dfs(dfs, 0);
  rep(i,n) ans[0] += dep[i];

  dfs(dfs, 1);
  coutv(ans);
 

  return 0;
}