NKOJ8663 彩色的树

· · 题解

_『果果』画了一棵\color{blue}{五}☆\color{orange}{彩}☆\color{red}{斑}☆\color{green}{斕}の树(｡>∀<｡)，用 1 \sim n 表示节点的编号，树上每个节点都有它对应的颜色 c_i。_

记 f(x,y) 表示 x 节点到 y 节点的最短路径上经过节点的颜色种类数量，求：

\sum_{i=1}^{n} \sum_{j=i+1}^{n} f(i,j)

看到熟悉的双层 \sum，不难想到要计算每一种值对答案的贡献。

但是如何拆分答案呢？

如果统计经过每一个点的路径会计算重复，不如直接对于每一种颜色分别计算经过这种颜色的路径数量。

但是依然不好直接计算。考虑正难则反，用总路径数量减去不经过这种颜色的路径数量。

因为树上的每一个点都是割点，所以剩下的路径只会在除去这些点外的连通块的内部。

不经过染色点的路径总数就是 \sum C_{size_i}^2，size_i 为每个连通块大小。

观察后发现，除了包含树根的连通块，其余的连通块都是在一个染色点每一个子节点的子树中删去染色点的子树得到的。因此我们可以在树上 dfs，在每一个连通块所在子树的根节点更新答案。

又因为每个节点的颜色是唯一的，所以以它作为连通块顶部上方一个节点可以更新的答案也是唯一的，一次 dfs 即可完成。

使用类似『Tarjan』的思想，在树上 dfs 过程中动态维护全局信息，并在每个节点被访问和回溯时进行处理。cnt_i 表示当前访问并回溯完的所有颜色为 i 的点的子树的并集大小。

第一次访问到某个点 u 时先记录当前的 cnt_{c_u}，再遍历 u 的所有子节点 v，用更新后的 cnt_{c_u}' 减去原本记录的 cnt_{c_u} 得到 v 点所在子树中颜色为 c_u 的节点的子树的并集大小，再用 v 的子树大小减去它，就可以得到以 v 节点作为顶端的不经过颜色 c_u 的连通块大小，以这个连通块的大小更新不经过颜色 c_u 的路径条数，并重置 cnt_{c_u}' 为原本记录的值，访问下一个子节点。

遍历完所有子节点后再将 cnt_{c_u}' 更新为原本记录的 cnt_{c_u} 加上 u 的子树大小。

dfs 过程完成后，用根节点的子树大小减去最终的 cnt_i，就可以得到上文所说的包含根节点的不经过颜色 i 的连通块大小。