官方题解只能做到 $O(n^4)$ / $O(n^3 \log n)$？玩的太差了。可以做到 $O(n^2\log^2 n)$ 时间、$O(n^2)$ 空间。关键是：不能再为每个结点显式计算完整的二维 DP 表，而要把重链内部的转移批量处理。 已有的二维树形背包优化可以达到 $O(n^3)$，但它仍然逐点处理这些二维表；下面的算法会绕开这一点。
状态与单点转移
题目统计的是可行的二元组 $(S,(c_u))$，不是染色方案数。因此，同一个颜色数量序列即使对应许多种染色，也只能计数一次。([QOJ][2])
记 $F_u[x,z]$ 表示：在 $u$ 的子树中，已经确定集合 $S$ 和所有结点的颜色数量，子树的颜色数量为 $x$，并且至少需要由 $u$ 的严格祖先提供 $z$ 种颜色的方案数。
这里的 $z$ 是最小需求，不是任意一次染色中实际出现的交集大小。如果一个方案的最小需求为 $z$，那么将一些原本不属于外部的颜色重新命名，可以让外部颜色数取遍 $z,z+1,\ldots,x$，同时保持子树中的颜色数量序列不变。这也是原题可行性判定所使用的贪心状态。([QOJ][3])
暂时允许根结点也选择属于 $S$。最终只统计 $F_0[x,0]$，根属于 $S$ 的方案至少需要一种外部颜色，自然不会计入答案。
设某个结点的儿子状态为 $(x_i,z_i)$。先把结点自身看成一个虚拟分量 $(1,1)$，记 $A=1+\sum_i x_i$、$B=1+\sum_i z_i$、$P=\max(1,\max_i x_i)$、$M=\max(1,\max_i z_i)$。
父结点的颜色数量可以取任意 $X\in[P,A]$。在处理父结点是否属于 $S$ 之前，最小外部颜色需求为 $Z=\max(M,B-A+X)$。
这是因为，将总颜色数从 $A$ 减少到 $X$，至多能让外部颜色数减少 $A-X$；同时，外部颜色数不能小于任何一个分量的需求。把各分量的外部颜色尽量重合，必要时将更多颜色重新命名为外部颜色，就能同时达到这两个下界。
随后有两个转移：父结点属于 $S$，得到 $(X,Z)$；父结点不属于 $S$，它自身提供一种颜色，得到 $(X,Z-1)$。由于已经加入 $(1,1)$，总有 $Z\ge1$。
固定儿子状态后，贡献的位置是一条折线：水平段为 $z=M$、$P\le x< A-B+M$，斜段为 $x-z=A-B$、$M\le z\le B$。原题解通过维护折线的几个端点，避免保存完整的四元组。([QOJ][3])
下面先说明代码如何高效实现一次转移，随后再消除逐点展开状态表的代价。
对若干分量维护三组统计：
- $\mathcal A[a,b]$：$\sum x_i=a,\sum z_i=b$ 的方案数。
- $\mathcal B[d,m]$：$\sum(x_i-z_i)=d,\max z_i\le m$ 的方案数。
- $\mathcal C[p,m]$：$\max x_i\le p,\max z_i\le m$ 的方案数。
$\mathcal A$ 使用二维普通卷积；$\mathcal B$ 对每个固定的 $m$ 做一维卷积；$\mathcal C$ 直接逐项相乘。这里使用的是最大值的前缀统计，所以后两种合并不需要枚举两个最大值。
令 $H[d,m]=\sum_{i=0}^{d}\mathcal B[i,m]$，令 $D[x,z]$ 为 $\mathcal A$ 的斜向前缀和，即 $D[x,z]=\mathcal A[x,z]+D[x-1,z-1]$。将折线的水平段、斜段分别差分，可以得到处理结点是否属于 $S$ 之前的分布：
$$ G[x,z]=\mathcal C[x,z]-\mathcal C[x,z-1] +H[x-z,z-1]-H[x-z-1,z]-D[x-1,z-1]. $$
负下标视为零。最后分别令 $F'[x,z]\mathrel{+}=G[x,z]$、$F'[x,z-1]\mathrel{+}=G[x,z]$。
这个公式还有一个重要性质：假如一个分量的 $x$ 范围是 $O(N)$、$z$ 范围是 $O(w)$，而其他分量的总大小为 $O(w)$，那么整次转移只需要处理 $O(N)\times O(w)$ 的矩形，时间为 $O(Nw\log N)$，不需要展开成 $N\times N$。
重链上的卷积分治
每个结点选择子树最大的儿子作为重儿子。对结点 $u$，定义 $w_u=1+\sum_{v\text{ 是轻儿子}}\operatorname{size}(v)$。一条重链上所有 $w_u$ 的和，恰好等于链头的子树大小。
先递归计算轻儿子的完整 DP，并把它们和虚拟分量 $(1,1)$ 合并。这样，在处理重链时，每个结点都成为一个作用在重儿子 DP 上的线性算子。
这个算子在远离 $z=0$ 的地方，实际上就是二维卷积。
设重儿子状态为 $(x,z)$，其余分量的总和为 $(a,b)$。如果 $z\ge w_u$，那么因为 $x\ge z$，重儿子的 $x,z$ 都不小于其余任何分量。因此，父结点颜色数为 $x+t$ 时，可以取任意 $0\le t\le a$，外部需求的变化为 $\max(0,t-a+b)$；父结点不属于 $S$ 时，再减一。
这个变化只依赖 $(a,b,t)$，不再依赖原来的 $(x,z)$。
用 $\xi,\eta$ 表示两个多项式变量。为了不存储负指数，定义平移后的卷积核：
$$ \widehat K_u(\xi,\eta) =(1+\eta)\sum_{a,b}\mathcal A_u[a,b] \sum_{t=0}^{a}\xi^t\eta^{\max(0,t-a+b)}. $$
这里的 $\mathcal A_u$ 只包含轻儿子和虚拟分量，不包含重儿子。实际转移是先乘以 $\widehat K_u$，再把 $\eta$ 指数整体减一。
$\widehat K_u$ 的两个次数分别不超过 $w_u$ 和 $w_u+1$。每个 $(a,b)$ 的贡献也是一段水平线和一段斜线，所以整个核可以在 $O(w_u^2)$ 时间内构造，不能再逐个枚举 $t$。
考虑重链上一段连续区间，结点数为 $h$，总权重为 $W=\sum w_u$。预处理区间核 $\widehat K_I=\prod_{u\in I}\widehat K_u$。
将输入 DP 分成两部分。
对于 $z\ge W$ 的部分，整段转移都可以使用卷积核。因为每经过一个结点，$z$ 最多减少一；在到达某个结点时，它仍然不小于尚未处理结点的总权重，因此一定满足当前结点的 $z\ge w_u$。同时，$x$ 从不减少，也不会触碰最大颜色数的边界。
所以，这部分只需要做一次二维卷积，乘以 $\widehat K_I$，最后把第二维整体左移 $h$。
对于 $z< W$ 的部分，不能直接套用整个区间核。但经过这段区间，$z$ 最多增加 $W$，因此它始终只会涉及高度 $O(W)$ 的区域。将区间继续分治，按原来的自下而上顺序处理即可。
注意，这不是把所有转移都当成可交换的卷积：只有确定不会触碰边界的部分才能批量卷积；低 $z$ 部分始终保持原转移顺序。
分治树按权重构造。选取一个结点，使其左侧和右侧的总权重都不超过当前区间的一半，将区间分成“左区间、这个单点、右区间”至多三个孩子。这样，分治深度为 $O(\log N)$，其中 $N$ 是链头子树大小。
处理一个区间时，算法如下：
- 将输入中 $z\ge W$ 的部分与区间核卷积。
- 将剩余部分依次交给几个孩子处理。
- 到达单点时，剩余输入满足 $z< w_u$，使用前面给出的精确转移。
- 将两部分答案相加。
初始输入是空重子树的状态 $F[0,0]=1$。整条重链处理完毕，才得到链头的完整 DP。
复杂度分析。 一条重链的总权重为 $N$。在某个总权重为 $W$ 的区间中，将低部分传给孩子时，其高度至多为 $2W$；横向范围始终为 $O(N)$。因此，该区间向孩子分发过程中涉及的卷积，总代价为 $O(NW\log N)$。
同一层的区间权重之和不超过 $N$，所以每层代价为 $O(N^2\log N)$；分治共有 $O(\log N)$ 层，得到每条重链 $O(N^2\log^2N)$。
单点精确转移的总代价为 $O(N\log N\sum_u w_u)=O(N^2\log N)$。轻儿子统计量和区间核使用同样的按权重分治合并，总预处理代价为 $O(N^2\log N)$。
最后，对所有重链链头求和。经过一条轻边，子树大小至少减半；轻边深度相同的链头，其子树两两不交。因此，所有链头子树大小的平方和满足 $\sum N^2\le n^2+n^2/2+n^2/4+\cdots=O(n^2)$。
所以总时间复杂度为 $O(n^2\log^2n)$。按权重构造分治树，也保证保存的卷积核及递归中的临时矩形总大小为 $O(n^2)$。