引理

在树结构中，任意三个不同的顶点 $u, v, w$ 之间一定会有一个唯一的中心交汇点（我们叫它 $c$）。 从 $c$ 出发，会像一个“丫”字形（Y型）一样，分别延伸出三条互不相交的路径到达 $u, v, w$（$c$ 也可以和 $u,v,w$ 中的某一个重合，此时就是一条线段上挂着另一个点，依然符合这个拓扑规律）。

基于这个引理，我们进一步统计：

• 位于 $c \to u$ 分支上的节点（不含 $c$），在 $p(u,v)$ 和 $p(w,u)$ 中各出现了一次，总共 2 次。
• 位于 $c \to v$ 分支上的节点，出现了 2 次。
• 位于 $c \to w$ 分支上的节点，出现了 2 次。
• 而那个唯一的中心交汇点 $c$，在三条路径里都出现了一次，总共出现了 3 次。

所以我们的问题可以转换为：枚举树上的每一个点 $u$ ，如果这个点的权重为完全平方数，我们便在它上面做统计。

问题又来了，怎么去做统计？

为了让 $u, v, w$ 的路径交汇在 $c$，它们的位置必须满足以下两种情况之一：

1. 三点都在外围： $u, v, w$ 分别处于 $c$ 的 三个不同的分支 中。
2. 一点在中心： $u, v, w$ 中的某一个点就是 $c$ 本身（假设 $w = c$），那么剩下的 $u$ 和 $v$ 必须处于 $c$ 的 两个不同的分支 中。

假设我们把 $c$ 剪掉后，它连着的各个分支的大小（节点数量）分别是 $S_1, S_2, S_3, \dots, S_k$。

• 情况 1 的组合数（从不同分支选 3 个点）：也就是在 $S_1 \dots S_k$ 中任选三个数字相乘，求总和。即 $\sum_{i < j < m} S_i S_j S_m$。
• 情况 2 的组合数（从不同分支选 2 个点，加上 $c$ 本身）：即 $\sum_{i < j} S_i S_j$。
• 有序是为了不重复统计。

但是对于每个点，我们不可能重新搜索一遍整棵树来得到它各个分支的大小。但是如果我知道了以 $c$ 为根的字数大小 $sz[c]$ ，我们只需要对子树维护好它的组合方式。这里使用动态前缀和的 trick。

假设我们当前正在遍历节点 $c$ 的各个分支，并且维护三个变量：

• $sum1$：目前已经遍历过的分支大小之和（选 1 个点的组合数）。
• $sum2$：目前已经遍历过的分支中，任选 2 个点所在不同分支的组合数。
• $sum3$：目前已经遍历过的分支中，任选 3 个点所在不同分支的组合数。

一开始，它们都是 0。当你拿到一个新的分支（假设大小为 $S$）时，按照从高到低的顺序更新：

1. 新增选 3 个点的组合： 就是在前方的旧分支里选 2 个点，再在当前新分支里选 1 个点。

   $$sum3 = sum3 + sum2 \times S$$

2. 新增选 2 个点的组合： 就是在前方的旧分支里选 1 个点，再在当前新分支里选 1 个点。

   $$sum2 = sum2 + sum1 \times S$$

3. 更新总大小：

   $$sum1 = sum1 + S$$

之后我们做 DFS 搜索即可。时间复杂度为 $O(n)$，非常优秀。