ABC133 F - Colorful Tree - かんプリンの学習記録

問題概要
思考の流れ
提出プログラム
感想

問題概要

${2\le N\le 10^5}$ 頂点の木が与えられる. 木の各辺には色 ${1\le c_i\le N-1}$ が塗られていて長さは ${1\le d_i\le 10^4}$ である. このとき, 以下の ${1\le Q\le 10^5}$ 個の問いに答えよ.

色 ${x}$ が塗られている辺の長さが ${1\le y\le 10^4}$ に変更されたと仮定したときの二頂点 ${u,v}$ 間の距離を求めよ. (辺の長さの変更はこれ以降の問いには影響しない)

思考の流れ

これっぽさがある. kanpurin.hatenablog.com

木上の二頂点間の距離は適当な頂点を根としてLCAを求めることで, 各頂点 ${v}$ から根までの距離 ${dist_v}$ を求めておくことですぐ求めることができる. しかし今回は辺の長さの変更があるので面倒.

${\downarrow}$

すべての色 ${c}$ と頂点 ${v}$ に関して根から ${v}$ までに通る色 ${c}$ に塗られた辺の数 ${cnt_{c,v}}$ とその長さの合計 ${len_{c,v}}$ がわかれば, 色 ${x}$ が塗られている辺の長さを ${y}$ に変更したときの根から頂点 ${v}$ までの距離は ${dist_{v}-cnt_{x,v}×y-len_{x,v}}$ となる.
よって質問に対する答えは ${(dist_{u}-cnt_{x,u}×y-len_{x,u})+(dist_{v}-cnt_{x,v}×y-len_{x,v})+2×(dist_{LCA(u,v)}-cnt_{x,LCA(u,v)}×y-len_{x,LCA(u,v)})}$ となる.

${\downarrow}$

各色各頂点に対して求めようとすると ${O(N^2)}$ かかり間に合わない.
そこで, DFSをする際に根から現在の頂点までの色 ${i}$ の辺の数とその長さの合計を保存する配列を(グローバルかなんかに)持つ. DFSで潜る際に辺の色と長さを配列に追加(足す)していってDFSで戻る際にその値を削除(引く)してやれば常に根から現在みている頂点までの必要な値が求まっていて, 更新も ${O(N)}$ 回で済んでいる.
質問を先読みして, どの頂点のどの色が必要なのかを見ておくことですべての値を保存しておく必要がなく高速に答えることができる.