整体思路

把原树以节点 $1$ 为根。对于任意一个合法点集，先观察其中深度最小的节点，也就是这个点集最靠近根的一层。

由于新图只连接原树上距离不超过 $2$ 的节点，所以这一层只可能有两种形态：

• 只有一个节点 $u$。
• 有多个节点，并且它们是同一个父亲 $u$ 的若干个同色儿子。

因此，可以按照一个合法点集的“最上层”唯一地统计它。

定义 $f_u$ 表示：包含节点 $u$，并且 $u$ 是点集中唯一最高点的同色连通点集数量。

接着考虑节点 $u$ 的一组同色儿子。它们两两在原树上的距离都是 $2$，因此在新图中两两相邻。设颜色为 $x$ 的儿子集合为 $C_x(u)$，那么从这些儿子中选择至少一个，并在每个被选儿子的子树中选择一个对应的 $f$ 状态，方案数为

$$g_{u,x} = \prod_{v\in C_x(u)}(1+f_v)-1.$$

这里的 $1$ 表示不选择儿子 $v$ 的子树，减去 $1$ 是排除所有儿子都不选的空集。

计算 $f_u$ 时，对于每个儿子 $v$，能够连接到 $u$ 的同色部分有两种：

• 如果 $c_v=c_u$，可以选择一个由 $f_v$ 统计的点集。
• 即使不选择 $v$，也可以选择 $v$ 的若干个颜色为 $c_u$ 的儿子，这部分由 $g_{v,c_u}$ 统计，因为这些节点与 $u$ 的距离都是 $2$。

所以有

$$f_u = \prod_{v\in \operatorname{son}(u)} \left( 1 + \mathbf 1_{c_v=c_u}f_v + g_{v,c_u} \right). $$

最后，所有最高层是某个节点同色儿子的合法点集，都被某个 $g_{u,x}$ 统计。只有最高点是根节点 $1$ 的点集无法被父亲统计，需要额外加入 $f_1$。

因此答案为

$f_1+\sum_{u=1}^{n}\sum_x g_{u,x}.$

整个过程自底向上进行，并将每个节点的儿子按颜色分组即可。