問題
次の 2 問は同じ設定ですが実行時間制限が異なります。
yukicoder No.3193 Submit Your Solution
yukicoder No.3194 Do Optimize Your Solution
作問: noya2 さん, 出題時の想定解法の提案 tatyam さん
概要
$N$ 頂点の木 $A$ , $B$ が与えられるので、 $\sum_{u=1}^N\sum_{v=1}^N\mathrm{dist}_{A}(u,v)\mathrm{dist}_{B}(u,v)$ の値を求めてください。
計算量は $O(N\log N)$ です。
解法
$S$ を $\lbrace 1,2,\ldots ,N\rbrace$ の部分集合、 $g$ を $N$ 以下の正整数とするとき、
$f(S,g)=\sum_{u\in S}\sum_{v\in S}\left(\mathrm{dist}_{A}(u,g)+\mathrm{dist}_{A}(v,g)\right)\mathrm{dist}_{B}(u,v)$
の値を求める問題を考えます。 $X_v=\mathrm{dist}_A(v,g)$ としてその値を既知とし、 $B$ 上で適当に根を取って根からの距離 $\mathrm{depth}$ と LCA を定義すれば、求める値は
$\sum_{u\in S}\sum_{v\in S}(X_u+X_v)(\mathrm{depth}_B(u)+\mathrm{depth}_B(v)-2\mathrm{depth}_B(\mathrm{lca}_B(u,v)))$
と表せます。 $B$ 上の木 DP (葉→根) を用いて $\mathrm{lca}_B(u,v)$ ごとに総和を求めるようにすれば、全体の総和を計算量 $O(N)$ で求めることができます。
$A$ の任意の頂点 $g$ をとると、 $uv$ 間単純パスが $g$ を通るような組 $(u,v)$ について $\mathrm{dist}_{A}(u,v)=\mathrm{dist}_{A}(u,g)+\mathrm{dist}_{A}(v,g)$ が成り立ちます。よって、上で求めた総和と元々求めたかった $\sum_{u=1}^N\sum_{v=1}^N\mathrm{dist}_{A}(u,v)\mathrm{dist}_{B}(u,v)$ との差は、 $A$ において $g$ をまたがないペア $(u,v)$ のみに生じます。そのようなペアについて足してしまった寄与は、 $A$ から $g$ を消したあとの連結成分の頂点集合を $S_1,S_2,\ldots ,S_k$ とすれば $\sum _{i=1}^kf(S_i,g)$ として同様に求められます。
これを再帰的に行います。つまり、 $g$ を取って $f(S,g)$ を計算して答えに加算し、分割された各部分 $V$ について $f(V,g)$ を計算して答えから引くということを、各 $V$ について再帰的に行います。
$g$ を毎回重心にとること（重心分解）により、 $S$ のサイズの総和は $O(N\log N)$ になります。再帰するごとに木 $B$ を $S$ の頂点を残して $O(|S|)$ 頂点に圧縮できるので、これにより木 DP の部分の計算量も $O(N\log N)$ となり、全体の計算量を $O(N\log N)$ とすることができます。
ここで $O(1)$ 時間の LCA を用いたこと、および圧縮後の木すべてについて頂点を DFS の行きがけ順に並べた列を得る必要があることに注意します。これは実装上の障壁なので、別の方法を使います。重心分解の分解過程で、部分木 $V$ が部分木 $V_1,V_2\,\ldots ,V_k$ に分解されるときに列 $V_1,V_2\,\ldots ,V_k$ をうまく binary splitting し、全体の分解過程を二分木で表せるようにしたとき、その二分木の高さを $O(\log N)$ に抑えることができます。 binary splitting にしたことにより、 $n$ 頂点の木から $m$ 頂点を保って圧縮した木を求めるアルゴリズムに $O(m)$ 時間の代わりに $O(n)$ 時間のものを用いても、全体の計算量が $O(N\log N)$ に保たれます。
実装
解答：
頂点集合 $V$ のときに選んだ $g$ を $g_V$ と表します。分割後に $f(g_V,V_i)$ を求めて引く部分と、直後の再帰で $f(g_{V_i},V_i)$ を求めて足す部分は、 $X_v=\mathrm{dist}_A(v,g_{V_i})-\mathrm{dist}_A(v,g_V)$ とすることで、一度に計算できます。
分割統治の流れで木を圧縮してゆくという方法は、オフラインクエリの分割統治で利用できます。(*1)
(*1) もっと定数倍が大変かと思いきや、そうでもなかったという感想があります。
補遺
似た問題に AGC047D – Twin Binary Trees があります。上で説明した方法を応用すれば、それぞれの木が完全二分木であるという制約を取り除いた状況での準線形時間の解法になります。(*2)
(*2) あのときのアレがこれです