Official

F - Tree Tree Tree Editorial by potato167


\(p\) を固定して、パスの長さが \(i\) となる頂点 \(u,v\;(u\lt v)\) の組の数を \(A_{i}\) としたとき、全ての \(q\) に対する「問題 potato」の答えの和はある数列 \(B=(B_{1},\dots ,B_{N-1})\) を用いて、 \(\sum_{i=0}^{N-1} A_{i}B_{i}\) と表せます。 よって、「問題 tomato」の答えは、\([x^i]f_{a}(x)\) を頂点 \(u,v\) のパスの長さが \(i\)\(p_{K}=a\) となるような \((p,u,v)\) の組の数とし、\(g(x)=\sum_{j=1}^{N-1}B_{j}x^{N-1-j}\) とすると、\([x^{N-1}] (f_{a}(x) g(x))\) と表せます。

まず \(B\) ですが、簡単に \(B_{j}=\frac{N! j}{j+1}\) であることがわかります。

\(f_{a}(x)\) はパスの種類で分解して、それぞれで求めます。

\(f_{a1}(x)\) : 頂点 \(u,v\) を結ぶパスに、頂点 \(a,K\) のいずれかが含まれていないもの

\(f_{a2}(x)\) : 頂点 \(u,v\) を結ぶパスに、頂点 \(a,K\) がどちらも含まれていて、\(u,v\) の LCA が \(a\) であるもの

\(f_{a3}(x)\) : 頂点 \(u,v\) を結ぶパスに、頂点 \(a,K\) がどちらも含まれていて、\(u,v\) の LCA が \(a\) でないもの

\(f_{a1}(x)\) について

\(f_{a1}(x)\) の条件に含まれているもののうち、LCA が \(1\leq i\leq N\) であるようなものを取り出した多項式を \(h_{i}(x)\) とすると、以下の式が成り立ちます。

\[h_{i}(x)=\frac{1}{2}\left(\prod_{2\leq j\leq i,j\neq K} j-1 \right)\prod_{i\lt j\leq N,j\neq K}((j-1)+2x)\]

ただし、\(h_{i}(x)\) の定数項はなんでも良いこととします。(結果に影響を与えないため) よって、 \(f_{a1}(x)=\sum_{}h_{i}(x)\) となります。これは分割統治 FFT で \(O(N\log^{2}(N))\) で求めることができます。

\(f_{a2}(x)\) について

\[f_{a2}(x)=x(a-1)!\prod_{j=a+1}^{K-1}((j-1)+x)\prod_{j=K+1}^{N}((j-1)+2x)\]

\(f_{a3}(x)\) について

\[f_{a3}(x)=x^{2}\sum_{i=1}^{a-1}\left((i-1)!\prod_{j=i+1}^{a-1}((j-1)+2x)\prod_{j=a+1}^{K-1}((j-1)+x)\right)\prod_{j=K+1}^{N}((j-1)+2x)\]

\(f_{a2}(x),f_{a3}(x)\) を制限時間内に陽に求めることはできませんが、 全ての \(a\) に対して、 \([x^{N-1}] (f_{a2}(x)g(x))\)\([x^{N-1}] (f_{a3}(x)g(x))\) を求めることは、必要のない項を適時削ることで同じく分割統治で \(O(N\log^{2}(N))\) で可能です。

\(f_{a3}\) についてのみ説明します。

任意の整数 \(0 \leq l\leq r\lt K\) に対して、 \(x\) の多項式 \(X_{l,r},Y_{l,r},Z_{l,r},W_{l,r},DP1_{l,r},DP2_{l,r}\) を以下のように定義します。

  • \(\displaystyle X_{l,r}=\prod_{i=l}^{r-1} \max(i,1)\)

  • \(\displaystyle Y_{l,r}=\prod_{i=l}^{r-1} (i+2x)\)

  • \(\displaystyle Z_{l,r}=\prod_{i=l}^{r-1} (i+x)\)

  • \(\displaystyle W_{l,r}=\sum_{j=l+1}^{r}X_{l,j}Y_{j,r}\)

  • \(\displaystyle DP1_{l,r}=g(x)x^{2}X_{0,l}Z_{r,K-1}\prod_{j=K}^{N-1}(j+2x)\)

  • \(\displaystyle DP2_{l,r}=g(x)x^{2}W_{0,l}Z_{r,K-1}\prod_{j=K}^{N-1}(j+2x)\)

\(g(x)f_{a3}(x)=DP2_{a-1,a}\) が成り立つので、全ての \(1\leq a\lt K\) について、 \([x^{N-1}] DP2_{i-1,i}\) を求めることができれば良いです。

任意の整数 \(0\leq l\lt m\lt r\leq K\) について、以下が成り立ちます。

\(DP1_{l,m}=DP1_{l,r}Z_{m,r}\)

\(DP1_{m,r}=DP1_{l,r}X_{l,m}\)

\(DP2_{l,m}=DP2_{l,r}Z_{m,r}\)

\(DP2_{l,r}=DP2_{l,r}Y_{l,m}+DP1_{l,r}W_{l,m}\)

\(DP1_{0,K-1},DP2_{0,K-1}\) は簡単に求まるので、そこから始めて、 \(m=\frac{l+r}{2}\) になるように \(m\) を選べば \(DP2_{a-1,a}\)\(O(\log K)\) 回の多項式掛け算で求めることができます。

\(DP2_{a-1,a}\) を求めるのに使う \(X,Y,Z,W\)\(DP1,DP2\) を計算する前に計算しておきます。

\(Y_{l,r},Z_{l,r},W_{l,r}\) は高々 \((r-l)\) 次式多項式なので、 \(r-l\) が小さくなればなるほどかけられる多項式の次数はどんどん小さくなります。

よって、事前計算で行う多項式の掛け算、足し算の、多項式の次数の和が \(O(K\log{K})\) になるため、時間計算量 \(O(K\log^{2}{K})\) で事前計算できます。

\(DP1_{l,r},DP2_{l,r}\) の計算も、 \([x^{N-1}] DP2_{a-1,a}\) を求めるには高々 \(O(r-l)\) 項しか残さなくていいため、\(DP1_{0,K-1}\) を求めるのに \(O(N\log^{2}{N})\) かかるのを除けば、 \(O(K\log^{2}K)\) で計算できます。

posted:
last update: