F - Tree Tree Tree Editorial by evima
If we fix \(p\) and let \(A_{i}\) denote the number of pairs of vertices \(u,v\;(u\lt v)\) with a path length of \(i\), the sum of the answers to the problem “potato” for all \(q\) can be expressed as \(\sum_{i=0}^{N-1} A_{i}B_{i}\) using a sequence \(B=(B_{1},\dots, B_{N-1})\). Therefore, the answer to the problem “tomato” can be expressed as \([x^{N-1}] (f_{a}(x) g(x))\), where \([x^i]f_{a}(x)\) is the number of triples \((p,u,v)\) such that the path length between vertices \(u\) and \(v\) is \(i\) and \(p_{K}=a\), and \(g(x)=\sum_{j=1}^{N-1}B_{j}x^{N-1-j}\).
First, it is easy to see that \(B_{j}=\frac{N! j}{j+1}\).
Let us decompose \(f_{a}(x)\) according to the types of paths and calculate each part separately.
\(f_{a1}(x)\): The number of paths connecting vertices \(u\) and \(v\) such that at least one of the vertices \(a\) and \(K\) is not contained.
\(f_{a2}(x)\): The number of paths connecting vertices \(u\) and \(v\) such that both vertices \(a\) and \(K\) are contained, and the LCA of \(u\) and \(v\) is \(a\).
\(f_{a3}(x)\): The number of paths connecting vertices \(u\) and \(v\) such that both vertices \(a\) and \(K\) are contained, and the LCA of \(u\) and \(v\) is not \(a\).
For \(f_{a1}(x)\)
Let \(h_{i}(x)\) be the polynomial obtained by extracting from \(f_{a1}(x)\) the terms where the LCA is \(1\leq i\leq N\), then the following equation holds:
\[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).\]
Here, the constant term of \(h_{i}(x)\) can be anything (because it does not affect the result). Therefore, \(f_{a1}(x)=\sum_{}h_{i}(x)\), which can be calculated in \(O(N\log^{2}(N))\) using divide-and-conquer FFT.
For \(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)\]
For \(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)\]
Although it is not possible to explicitly calculate \(f_{a2}(x)\) and \(f_{a3}(x)\) within the time limit, it is possible to calculate \([x^{N-1}] (f_{a2}(x)g(x))\) and \([x^{N-1}] (f_{a3}(x)g(x))\) for all \(a\) in \(O(N\log^{2}(N))\) using divide-and-conquer by appropriately removing unnecessary terms.
Let us explain only \(f_{a3}\).
For any integers \(0 \leq l\leq r\lt K\), define the polynomials \(X_{l,r},Y_{l,r},Z_{l,r},W_{l,r},DP1_{l,r},DP2_{l,r}\) of \(x\) as follows:
\(\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)\)
Since \(g(x)f_{a3}(x)=DP2_{a-1,a}\) holds, it is sufficient to find \([x^{N-1}] DP2_{i-1,i}\) for all \(1\leq a\lt K\).
For any integers \(0\leq l\lt m\lt r\leq K\), the following holds:
\(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}\)
Since \(DP1_{0,K-1}\) and \(DP2_{0,K-1}\) can be easily found, starting from there and choosing \(m\) to be \(m=\frac{l+r}{2}\), one can find \(DP2_{a-1,a}\) in \(O(\log K)\) polynomial multiplications.
The \(X,Y,Z,W\) used to find \(DP2_{a-1,a}\) should be calculated before calculating \(DP1\) and \(DP2\).
Since \(Y_{l,r},Z_{l,r},W_{l,r}\) are at most \((r-l)\)-degree polynomials, as \(r-l\) decreases, the degrees of the multiplied polynomials also decrease.
Therefore, the polynomials in the polynomial multiplication and addition performed in the precomputation have a total degree of \(O(K\log{K})\), so the precomputation can be done in time complexity \(O(K\log^{2}{K})\).
The calculation of \(DP1_{l,r}\) and \(DP2_{l,r}\) can also be done in \(O(K\log^{2}K)\) time, except for the \(O(N\log^{2}{N})\) time required to find \(DP1_{0,K-1}\), because only at most \(O(r-l)\) terms need to remain to find \([x^{N-1}] DP2_{a-1,a}\).
posted:
last update: