First, let’s think of assigning a polynomial \(\mathbf{F}_i\) for each vertex \(i\), such that the coefficient of \(x^k\) is the number of good vertex sets with exactly \(k\) vertices for the subtree at vertex \(i\).

We can see that \(\mathbf{F}_i= x + \prod_{j \text{ is a child of } i} \mathbf{F}_j\) . (If vertex \(i\) has no children, the product is considered to be \(1\), the “empty product”), as we can either choose to color only vertex \(i\), or pick some coloring of the children’s subtrees and not color vertex \(i\). However, implementing this naively can result in \(O(N^2)\) time or worse, even with FFT and the trick of using a queue / divide-and-conquer to multiply several polynomials.

(To see why, picture a “centipede” tree that resembles the diagram below)

  1
 / \
2   3
   / \
  4   5
     / \
    6   7
       / \
      8  ...

Instead, we consider the heavy-light decomposition (HLD) of the tree. For each non-leaf vertex, we shall pick the child with the largest subtree (in terms of number of vertices), breaking ties arbitrarily, to be the heavy child, the edge toward that child becoming a heavy edge. Other children and edges remain light.

It can be shown that the path from the root to any vertex does not cross more than \(O(\log N)\) light edges. This is because whenever we descend down a light edge, the subtree size of that child must be strictly less than half the subtree size of the parent. Below is a visualization of the HLD of a tree, where red solid edges are heavy and blue dashed edges are light:

How does this help us? We can try to defer the calculation of \(\mathbf{F}_i\) until we reach either a light edge or we get to the root of the tree (vertex \(1\)). Let us define a function \(\mathbf{H}(i)\) that returns an array of polynomials (implemented by a C++ vector or a Java ArrayList, or similar analogues) according to the following procedure:

• If \(i\) is a leaf vertex, simply return an array with a single polynomial \(1\).
• Otherwise, call \(\mathbf{H}(h)\), where \(h\) is the heavy child of \(i\). Also calculate \(\prod_{j \text{ is a light child of } i} \mathbf{F}_j\). Append that product to the \(\mathbf{H}(h)\) array and return it.

We also need to think about how to calculate \(\mathbf{F}_i\) from the polynomial array \(\mathbf{H}(i)\). Let that array be \([\mathbf{H}_1, \mathbf{H}_2, , ..., \mathbf{H}_m]\). Note that the length of this array is the same as the length of the “heavy path” rooted at \(i\), we can represent that path going upward as an array: \([p_1, p_2, , ..., p_m = i]\). It can be seen that \(\mathbf{F}_{p_{1}} = x + \mathbf{H}_1\), and \(\mathbf{F}_{p_{j}} = x + \mathbf{H}_j\mathbf{F}_{p_{j-1}}\). However we still run into the problem of this taking \(O(m^2)\) or worse.

Instead, let’s consider this recurrence as an array of affine functions of the form \(f(y) = \mathbf{A}x + \mathbf{B}y\), where \(\mathbf{A}\), \(\mathbf{B}\) and the result of invoking \(f(y)\) are all polynomials over \(x\). We start with the array \([f_1(y), f_2(y), ..., f_m(y)] = [\mathbf{1}x + \mathbf{H}_1y, \mathbf{1}x +\mathbf{H}_2y, , ..., \mathbf{1}x + \mathbf{H}_my]\), it can be seen that \(\mathbf{F}_i = f_m(...(f_2(f_1(\mathbf{1}))))\) . We can compose two affine functions \(f(y) = \mathbf{A}x + \mathbf{B}y\) and \(g(y) = \mathbf{C}x + \mathbf{D}y\) via substituting \(y\) in the latter function with the former function: \(g(f(y)) = \mathbf{C}x + \mathbf{D}(\mathbf{A}x + \mathbf{B}y) = (\mathbf{AD} + \mathbf{C})x + \mathbf{BD}y\), this produces another affine function and takes time \(O(n \log n)\) with respect to the result length. As composition is associative, we can use divide-and-conquer to compose the entire array of affine functions in \(O(s \log s \log m)\) time, where \(s\) is the subtree size of vertex \(i\). We can then obtain \(\mathbf{F}_i\) by substituting \(y = 1\) (i.e. \(\mathbf{A}x + \mathbf{B}\)) in the final composed affine function.

To consider the final time complexity, we note that evaluating \(\mathbf{F}_i\) is only done once per light edge and once at the root, and each time it takes \(O(s \log s \log m)\) time, which we shall simplify to \(O(s \log^2 s )\). Calling the helper function \(\mathbf{H}(i)\) also takes \(O(l \log^2 l)\) time (if we also use divide-and-conquer for the product), where \(l\) is the sum of the subtree sizes of \(i\)’s light children. Now note that each vertex can only contribute to \(s\) and \(l\) up to \(O(\log N)\) times (proportional to the number of light edges from it to the root) thus the final time complexity is bounded by \(O(N\log^3 N)\).

This editorial is written via interpretation of Nyaan’s solution (C++). Note that there are two styles of divide-and-conquer in this code: “top-down” using recursive splitting (used for the affine functions), and “bottom-up” via repeated compression of the array (used for the product of polynomials from light children), either would work for both cases. (Note that the queue method that’s sometimes used elsewhere for product of polynomials would not work for the affine functions, as the composition operation isn’t commutative.)