F - path pass i Editorial by maple116

重心分解による解法

 条件をみたすパスの数え上げにはしばしば重心分解が有効であり、この問題も例外ではありません。以下では重心分解による解法を紹介します。なお重心分解の基本的な原理および実装などについては各種資料をご覧ください。

 唐突ですが、まずは以下の小問題を考えましょう。

  • 各頂点が \(N\) 色のいずれかで塗られた \(N\) 頂点の根付き木に対し、色 \(k\) が塗られている頂点を一度以上通り、かつ根を通るような(単純)パスの数を、\(k=1,2,\cdots,N\) に対してそれぞれ求めよ。

 なお以下では簡単のため根の色については例外視します(適切な処理を行ってください)。

 頂点 \(v\) を根とし、その子孫をすべて含むような部分木を \(T_{v}\) とおきます. ただし頂点 \(0\) が元の木の根であるとします(すなわち \(T_{0}\) は元の木全体)。さらにDFSによって各頂点に対し \(T_{v}\) に含まれる頂点の数 \(|T_{v}|\) をボトムアップに前計算しておきます。

 ここで頂点 \(u,v\) に対し、二項関係 \(\leq\) を「\(u\)\(v\) の祖先であるとき \(u\leq v\)」として定め(ただし自身も祖先とする)、\(N\) 個の頂点全体を半順序集合とみなします。さらにその極小元であって色 \(k\) で塗られたもの全体に対し \(|T_{v}|\) の総和を \(S_{k}\) とします。これはDFSによって計算できます。

 また根の次数を \(d\) とし、根の子孫が頂点 \(1,2,\cdots,d\) であるとします。このとき \(i=1,\cdots,d\) に対し各 \(T_{i}\) において上の \(S_{k}\) に相当する値 \(S'_{i,k}\) を同様に計算しておきます。

 結論から言うと、このとき求める答え \(ans_{k}\) は以下のように計算できます。

\[ ans_{k} = \sum_{i=1}^{d} \left(S'_{i,k}(|T_{0}|-|T_{i}|)+\frac{{S'}_{i,k}^{2}}{2} \right)-\frac{S_{k}^{2}}{2} \]

 これで正しい答えが出る理由を確認しましょう。同じパスを両方向から加算するとして \(2ans_{k}\) を求めることを考えると簡明です。ここで各色で塗られた頂点について特に極小元であるようなものだけを考慮すれば良いことに注意してください。さらに根を通るパスはある色の極小元を高々二つしか通りません。

 上式の第一項では色 \(k\) の極小元を二つ通るものが余分にカウントされているため、それを第三項で差し引いています。しかしこのとき根を通らないようなものまで不要に差し引いているため、これを第二項でさらに調整しています。

 なおここで数列 \(a_{1},\cdots,a_{n}\) に対する以下の公式を背景としていることに留意してください。

\[ \sum_{i\neq j}a_{i}a_{j}=\left(\sum_{i}a_{i}\right)^{2}-\sum_{i}a_{i}^{2} \]

 上式を愚直に計算すると間に合いませんが、各段階において更新が必要な色の数は限られているため、これを適切に管理することで小問題は \(O(N)\) で解けました。

 さらにこれを分割された各部分木においてにおいて重心を根として適用することで、重心分解の原理より全体で \(O(N\log N)\) で元問題が解け、これは十分高速です。

 なお実装方針によって重めの定数倍がかかり得るため、言語によっては実行時間制限をクリアできない可能性があります。

posted:
last update: