ABC133 F. Colorful Tree

問題

N個の頂点からなる木がある。
各辺iは、頂点a_iとb_iをつないでおり、1～N-1の整数で表される色c_iと、辺の長さd_iが割り当てられている。
このとき、以下のQ個の問いに答えよ。
「問いj: 色x_jのすべての辺の長さをy_jに変更したと仮定して、二頂点u_j, v_jの距離を求めよ」

(注意：各問いは独立で、変更は他の問いに影響しない)

制約

• 2 <= N <= 10^5
• 1 <= Q <= 10^5
• 1 <= c_i <= N-1
• 1 <= d_i <= 10^4

解法

まず、単純に2頂点間の距離を答えるだけの場合は、根付き木にして、LCAを求めて、「depth[u_j] + depth[v_j] - 2 * depth[lca(u_j, v_j)]」で求められる。

問いは、色x_jの辺の重みをy_jにするので、根付き木において、ある範囲内における色x_jの「辺の重みの合計」や「辺の数の合計」がわかれば、

答え＝「(u_jとv_jの距離) - (u_jとv_jの間の色x_jの辺の重み合計) + y_j * (u_jとv_jの間の色x_jの辺の変数合計)」

で求められる。
なので、基本的に、「根付き木について、根からある頂点までの間にある色がxの辺の重みや個数」を求めることが本質的な問題になっている。

1を根として、根付き木を考え、辺についてのEuler Tourを作ると上記のようになる。
行きがけは重みを＋、帰りがけは重みをーにして累積和を取ると、ある頂点vの根からの距離は、頂点vの根方向への辺eの重みの累積和を見ればO(1)で求められる。
(例：頂点5は、根方向への辺がe4で、それに対応するeuler tourのインデクスはi=3なので、その重みの累積和は50、とわかる)
辺の数についても同様に求められる。

ある色xについて考えると、この累積和のように、範囲内にある、色がxのものだけについて注目して累積和を求められばよいことがわかる。

なので、Euler Tourを各色ごとに分解して累積和を計算しなおしたものを考える。

ある頂点vについて、根方向への辺eに対応するeuler tourのインデクスがiだった場合を考えると、色がxのものだけに注目したい場合、その色だけ取り出した表の中でi以下となるインデクスを二分探索で求めてあげれば、ある色の辺だけについて重みの合計を求められる。
辺の数についても同様に求められる。

したがって、1つの問に対して、O(logN)程度で答えられるので、間に合う。

Euler Tourを使ったLCA

Euler Tourを作る際に、根からの深さを一緒に求めておく。
これを、RMQなどで、あるEuler Tourの範囲での最小の高さになっている辺(頂点)を求めてあげればよい。