AtCoder Beginner Contest 133 F - Colorful Tree

F - Colorful Tree

解法

木の2点間距離たくさん→LCA

長さの変更についていったん無視すると、この問題は「木の2点間距離をたくさん求めてください」というものです。これはLCA（最小共通祖先）というものを用いると効率的に求めることができます。

根付き木における2頂点のLCAとは、それらの共通祖先のうち最もその2頂点から近い頂点のことです。2頂点からそれぞれ根までのパスを考えた時に合流するところ、とイメージしてもよいです。

これを用いると、2点間の距離は以下のように求めることができます。

前計算
- 適当に根を決める。
- あらかじめ、根から全点への距離を求めておく。根と頂点 $v$ の距離を $d_{v}$ と表す。
- 2点のLCAを効率的に求めるための前準備をしておく。
頂点の距離の計算
- 頂点 $u, v$ のLCAを求める。これを $l$ とする。
- 頂点 $u, v$ 間の距離を、 $d_{u} + d_{v} - 2d_{l}$ と求める。

f:id:betrue12:20190707233844p:plain

LCAを効率的に求める方法としてはダブリングなどがあります。（仕組みや実装例は蟻本などを参照してください）

各クエリで辺の値が変わるときにもこのテクニックは使えます。つまり各クエリ $(x, y, u, v)$ に対しては、 $l$ を $u, v$ のLCAとして「色 $x$ の辺の長さを全て $y$ に変えた時の、ある1頂点と根との距離を求める」ことを $u, v, l$ の3頂点について求めればよいことが分かりました。これ以降は根との距離を上手く求めることを考えていきます。

長さ変更の影響を考える

色 $x$ の辺の長さを $y$ に変えた時の、頂点 $v$ から根までの距離は

頂点から根までの経路の…
- 初期状態における距離 $D_{v}$
- 経路中にある色 $x$ の辺の長さ合計 $d_{v, x}$
- 経路中にある色 $x$ の辺の本数 $n_{v, x}$

これらを用いて $D_{v} - d_{v, x} + yn_{v, x}$ と求めることができます。

これらの値のうち、 $D_{v}$ を全点について求めることはDFSなどで根から順に計算していけば $O(N)$ で可能です。一方色に依存する $d_{v, x}$ や $n_{v, x}$ は、色ごとに計算すれば求めること自体は $D_{v}$ と同様にDFSでできるものの、全ての頂点×色の組み合わせが $O(N^{2})$ 個あるので全部求めるとTLE/MLEなどになってしまいます。

クエリ先読み！

そこでまず、DFSをしながら「今見ている頂点についての $d_{v, x}, n_{v, x}$ はいくつか？」という値を全ての色 $x$ について管理していくことを考えましょう。DFSで1歩進む・戻るときに通る辺は1本なので、その辺の色についての情報だけを変更すればよいです。

これを全ての頂点×色について記憶するとTLE/MLEになりますが、クエリは事前に分かっているので必要なところだけを記憶しておくことができます。各クエリに対して必要な頂点と色のペアは3個なので、必要になるのは全部で最大でも $3Q$ 個のペアです。あらかじめこれを列挙して頂点ごとに分類しておき、DFSしながら各頂点を通る時に必要な色についての値だけ記憶しておくようにすれば、DFSと差分更新に $O(N)$ 、値の記憶に（mapなどを使うことを想定して） $O(Q\log Q)$ の時間計算量で前計算をすることができます。