First of all, if you are interested in the solution of the problem itself, we recommend you to start using from “editorial for this problem” in the last part of this page.

First, we describe a technique called Heavy-Light Decomposition which we need to solve this problem.

Heavy-Light Decomposition

The Heavy-Light Decomposition (HLD) is a technique to divide a rooted tree into several paths.
Converting a tree to several paths with HLD enables us to apply algorithms applicable to sequences to trees too.

Algorithm

• Note: the algorithm we describe below is precisely not HLD, but rather considered to be Centroid Path Decomposition, but most Japanese competitive programmers know this algorithm as HLD, and is easy to implement, so we introduce that.

Consider an \(N\)-vertex rooted tree rooted at vertex \(1\).
First, we precompute the sizes of all subtrees. Then, for each vertex \(v\), we classify each edge from \(v\) to its child into either a heavy edge or a light edge as follows:

• Choose one child of \(v\) with the largest subtree and connect it and \(v\) with a heavy edge. (If there are multiple of them, any of them can be chosen.)
• The other children and \(v\) are connected with light edges.

Here, we will denote the children connected with a heavy or light edge by a heavy or light child, respectively.
Then, each vertex has at most one heavy edge towards the children. Here, the set of vertices connected with heavy paths forms a path graph. We call it a heavy path.
By the operations above, we have now decomposed the tree into a set of heavy paths and of light edges.

• An example decomposition (Red Bold lines are heave edges and blue are light edges):

Property

A HLD has the following property:

• The path from any vertex from any vertex \(v\) to the root \(1\) can be decomposed into \(\mathrm{O}(\log n)\) heavy paths and light edges.

We assert the validity. Consider traveling from \(v\) towards the root. Suppose that we have traveled from vertex \(s\) to vertex \(t\) via a light edge. Then it holds that

• (the size of subtree of \(s\)) \(\geq\) \(2 \times\) (the size of subtree of \(t\))

because \(s\) has a child that forms a subtree larger than the size of subtree of \(t\) as a heavy child. Thus, you go through light edges at most \(\lceil \log_2 N \rceil\) times; since heavy edges forms path graphs, the claim is justified.

Similarly, a sequence of vertices contained in the path from any vertex \(v\) to the root can be expressed as \(\mathrm{O}(\log N)\) sequences of vertices on heavy paths. Moreover, with LCA (Least Common Ancestor), any path between two vertices \(u-v\) can be decomposed into \(\mathrm{O}(\log N)\) heavy paths.

Therefore, by maintaining the heavy paths as a sequence beforehand, an operation on any path can be boiled down to operations on \(\mathrm{O}(\log N)\) edges.

Using preorder for implementation

The easiest way we think is the best is to use together with preorder.

First, calculate the size of subtree of each vertex to split them itno heavy edges and light edges. One of the easiest way is to store the heavy children of \(c\) into the following adjacent list g:

vector<vector<int>> g; // The adjacent list for the rooted tree
int sub[N_MAX + 3];    // The size of the subtree rooted at i

int dfs_hld(int c) {
  sub[c] = 1;
  for(auto& d : c) {
    sub[c] += dfs_hld(d);
    // Swap if the currently focused subtree is larger than that rooted at $g[c][0]$
    if(sub[d] > sub[g[c][0]]) swap(d, g[c][0]);
  }
  return sub[c];
}

Then we find the preorder of g by DFS. We also find the (endpoint on the root side of the heavy path containing \(n\)). (We store it in head in the code below.)

int order[N_MAX + 3]; // the order of i occuring in the preorder
int head[N_MAX + 3];  // the endpoint of the heavy path containing i
void dfs2(int c, int &i) {
  order[c] = i++;
  for(auto& d : c) {
    head[d] = (g[c][0] == d ? head[c] : d);
    dfs2(d, i);
  }
}

Then the following property of heavy paths and preorder holds:

• Any heavy path is a contiguous sequence on the preorder.

This is because for all \(n\), the head of the adjacent list g[n] is a heavy child. Thus, we can exploit the preorder to manage heavy paths easily.

Also, we may use the precalculated head to enumerate the heavy paths that is contained in the path from any vertex \(v\) to the root. (With a tweak on the precalculation, we may also find the sequence of heavy paths containing any path \(u-v\). Do consider how can we do it.)

Therefore, by sorting the data in, for example, the preorder, and managing them on a data structure like a segment tree, one can convert a retrieval or updating query on a path on the tree to that on \(\mathrm{O}(\log N)\) sequences.

See also: article in CodeForces

Another typical usage of HLD in contests is to apply Weighted-Union Heuristic on a given tree. This problem uses an algorithm that is classified to that kind.

Editorial for this problem

First of all, this problem can be solved in a total of \(\mathrm{O}(N^2)\) tie with a quadratic-order tree DP (Dynamic Programming). The transitions seems welcoming optimization with FFT (Fast-Fourier Transform), but it is difficult to get AC (Accepted) on a centipede graph, which would cost \(\Theta(N^2)\) time.

Here, we denote by \( f_n(x)\) the generating function of the answer for the subtree rooted at \(n\).

First, let us consider the following subproblem:

• There is an \(N\)-vertex tree rooted at \(1\). \(1-2-3-\dots-M\) is a path on this tree. We have found \(f_v(x)\) for all vertex \(v\) \((M \lt v \leq N)\) that is not on the path.

This problem can be solved in a total of \(\mathrm{O}(N \log^2 N)\) time.
Here is an simple explanation. (Although not easy, it is a moderate difficulty of the practice of two-log problems, so consider it if you don’t get it.)
Let \(g_n(x)\) define as follows:

\[g_n(x) = \prod_{P_v = n, v \gt M} f_v(x).\]

Then, with a proper observation, we can see the following relation between \(f_1(x)\) and \(g_{\ast}(x)\), where we can apply divide-and-conquer:

\[f_1(x) = \sum_{i=1}^{M+1}\left( (\text{1 if }i = M+1\text{ else }x) \prod_{j=1}^{i-1} g_j(x) \right).\]

Now let us consider the general case. First, compute the sizes of subtrees to classify edges to heavy/light ones. Then, consider evaluating DP from leaves using DFS.

• By the way, in some languages it may be easy to implement scanning \(n=N, N-1, \dots, 2, 1\) instead of DFS in this problem (because the order written above is guaranteed to be the reverse of the topological order, since \(P_i \lt i\)).

Let \(\text{heavy}_n\) be the heavy child of vertex \(n\). For each vertex \(n\), just as we replaced the “path” in the subproblem with a “heavy path,” we define \(g_n(x)\) by the product of \(f_{\ast}(x)\) for the light children of \(n\), as follows:

\[g_n(x) = \prod_{P_v = n, v \neq \text{heavy}_n} f_v(x).\]

For a vertex \(n\), let \(X_{n, 1}, X_{n, 2}, \dots, X_{n,k}\) be the sequence of vertices on the heavy path starting from \(n\) towards its children. Then, we consider the sequence \(F_n = (g_{X_{n,1}}, \dots, g_{X_{n,k}})\) consisting of \(g_{\ast}(x)\) of each vertex on the sequence of vertices. We manage this \(F_n\) instead of \(f_n\) during the DP.

First, if it doesn’t have a child, \(F_n = (1)\). We now assume that it has a child.
First, we obtain \(F_{\text{heavy}_n}\) that we have already computed in an \(\mathrm{O}(1)\) time (via a pointer or std::move) to set \(F_n \gets F_{\text{heavy}_n}\).

Then we merge the DP result of light children of \(n\) to obtain \(g_n(x)\). \(f_v(x)\), which is required to compute \(g_n(x)\), can be found by converting \(F_v\) to \(f_v(x)\), just as the divide-and-conquer in the subproblem. Finally, we append \(g_n(x)\) to \(F_n\).
By computing this from the leaves with DFS, we can find \(F_1\).

Let us consider the complexity. The whole except for the following two parts can be processed in a total of \(\mathrm{O}(N)\), so we will consider:

• the part when converting \(F_n\) into \(f_n(x)\);
• the part when merging the results for light edges and computing \(g_n(x)\).

Converting \(F_n\) into \(f_n(x)\) costs \(\mathrm{O}(S_n \log^3 S_n)\) time, where \(S_n\) is the size of the subtree rooted at \(n\). Also, we need to convert \(F_n\) into \(f_n(x)\) only for light children and Vertex \(1\). Here, by the property of HLD,

\[\sum_{\text{heavy}_{P_v} \neq v} S_v = \mathrm{O}(N \log N).\]

So we can assert that the complexity of this whole DP is bounded by \(\math\)\mathrm{O}(N \log^3 N)\(rm{O}(N \log^3 N)\). Also, merging the results for light edges is also bounded by \(\mathrm{O}(N \log^3 N)\). Therefore, the complexity of the overall algorithm is bounded by \(\mathrm{O}(N \log^3 N)\) as well.

Considering the constant factor, one of the \(\log N\) in the complexity is due to the property of light edges, and another is the logarithm of lengths of heavy path, which cannot reach the upper bound at the same time, so the constant factor becomes lighter due to this. Therefore, the algorithm can be considered working fast enough.

Sample code (C++, 454 ms)

In the sample code above, we use a function in c++ called std::move when passing data via a heavy edge. (In C++11 and later, the return value of a function is automatically moved when a certain condition is met. For more details, please refer to the chapter “Automatic move from local variables and parameters” in the article of return statement in cppreference.)

P.S. Somebody told me a solution with a better complexity, so we will introduce it here.
It seems we can show that the results of light edges can be merged by “choose two polynomials with the smallest degree and replace them with the product,” where the complexity is bounded by \(\mathrm{O}(N \log ^2 N)\). Also, determining the values for heavy paths with divide-and-conquer can also be done with one less \(\log\), and as a result, it seems to be able to be computable in a total of \(\mathrm{O}(N \log^2 N)\) time. (We don’t know the details.):

• See also: article by Nachia (Japanese) article by yosupo (Japanese)