Official

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: