Subtask 1(5 points)
直接暴力建树,写基环树 dp,或者发现答案只与环长有关都可以拿到这档部分分,复杂度 $\Theta(n^3 A)$。
Subtask 2(30 points)
由Subtask1(或者其它手段),我们发现答案只与环长有关,具体来讲考虑确定了环的方案后,树上的点只有一个限制就是不与父亲相同,所以树上点的方案就是直接乘 $k-1$。
而这棵基环树的环长就是 $A \cdot (d(s,t)+1)$,所以可以枚举 $s,t$ 求距离,求完距离之后可以用矩阵快速幂求出环长为 $len$ 的答案,注意这里不同的环长最多只有 $\Theta(n)$ 种,所以把环长相同的放到一起算即可。
如何矩阵快速幂?我们设 $[\begin{matrix} f_{i,0} & f_{i,1} \end{matrix}]$ 表示链长为 $i$ 的链尾的颜色与链头是否相同的染色方案数,每当链长度加一就相当于我们给 $[\begin{matrix} f_{i,0} & f_{i,1} \end{matrix}]$ 乘上转移矩阵 $\mathbf M = \left[\begin{matrix} 0 & k - 1 \\ 1 & k - 2 \end{matrix}\right]$ 得到 $[\begin{matrix} f_{i+1,0} & f_{i+1,1} \end{matrix}]$,所以递推矩阵应该是这样:$[\begin{matrix} 1,0 \end{matrix}] \times \mathbf M^n$,不同的长度只有 $\Theta(n)$ 种,每次暴力计算,复杂度 $\Theta(n^2+b^3n \log A)$,其中 $b$ 为转移矩阵的阶。
Subtask 3(25 points)
因为要求出树上任意两点间的距离,对每个距离都要求出数量,所以我们会很自然的想到树上点分治+FFT/NTT,但 是这个模数很黑,它只有 $14$ 次单位根 $\omega_{14}$,所以 NTT 是不可行的,而 FFT 精度会出问题(这个点 比较小也许不会出问题),所以只能 MTT,总复杂度 $\Theta(n \log^2 n + b^3 n \log A)$。
Subtask 4(40 points)
因为我们要求的并不是距离而是矩阵,所以考虑把求距离和矩阵乘法两部分合到一起做。
设 $1$ 为根,,设 $\mathbf G_x$ 表示 $x$ 子树中的所有点到 $x$ 的矩阵之和,即:
$$\mathbf G_x=\sum_{y \in \mathrm{subtree}(x)} \mathbf M^{A \cdot (d(x,y)+1)} (k-1)^{A(n-d(x,y)-1)}$$
那么所有 $s=x,t\in \mathrm{subtree}(x)$ 的贡献我们就一并统计了,此时 $G_1$ 就是 $s=1$ 的答案。$G_x$ 的转移式也很好写出:
$$\mathbf G_x=\frac{1}{(k-1)^A}\left(\sum_{y \in \mathrm{son}(x)} \mathbf G_y + \mathbf I \times (k-1)^{An}\right) \times \mathbf M^A$$
其余点的贡献可以直接换根 dp 求出。时间复杂度 $\Theta(b^3n)$。