Colorful Tree

题目链接： HDU 2017 Colorful Tree 。

题目描述：

有 $n$ 个点的树，每个点有一个颜色 $a_i$ ，求所有 $(u,v)$ 满足 $1\le u<v\le n$ 的点对的路径上的不同颜色的和。

$n\le 2\times10^5,1\le a_i\le n$ 。

考虑对于枚举每一个颜色并计算它带来的贡献。对于枚举的颜色 $c$ ，我们应该视所有 $a_i=c$ 的 $i$ 为关键点，那么这个颜色 $c$ 的贡献就是所有经过了关键点的路径。

这个并不好算，我们考虑把它容斥为不经过任意一个关键点的路径。

此时问题就明了了：将所有的关键点去掉后的连通块的大小的平方和。

由于有关键点，我们对所有的关键点建立虚树。为了方便强行钦定 $0$ 为根，连边 $(0,1)$ ，并设 $0$ 为关键点。

接下来就我们不妨自下而上考虑：对于一个关键点 $u$ ，遍历它在原树上的儿子 $v$ ，并计算 $v$ 的子树权值和（即去掉 $u$ 之后的 $v$ 这一个连通块大小）的平方，然后将 $u$ 整个子树的权值全部置为 $0$ 。初始时所有节点均置为 $1$ 。

复杂度容易说明：建立虚树的复杂度 $\mathcal{O}(n\log n)$ ，之后每一个节点都只会出现在一个虚树中并遍历一次它的儿子，执行一遍线段树操作，于是这部分复杂度 $\mathcal{O}(n\log n)$ 。总复杂度就是 $\mathcal{O}(n\log n)$ ，可以通过。

代码十分简单，不放了。

PS: 这里虚树的作用其实是为了保证自下而上的遍历，所以也可以不用虚树，改为将所有的关键点按深度排序处理。

直接枚举所有路径并统计颜色种类的复杂度是 $O(n^2)$，对于 $n \le 2 \times 10^5$ 的数据范围会超时。

我们需要转换思维，从“每条路径有多少种颜色”转变为“每种颜色对总答案的贡献是多少”。

假设树上一共有 $n$ 个节点，那么树上的点对（路径）总数为 $\frac{n(n-1)}{2}$。

对于任意一种颜色 $c$，它对答案的贡献就等于包含至少一个颜色 $c$ 节点的路径数量。

直接求包含颜色 $c$ 的路径比较困难，我们可以使用补集思想：包含颜色 $c$ 的路径数 = 总路径数 - 完全不包含颜色 $c$ 的路径数。

如果我们把树上所有颜色为 $c$ 的节点删除，整棵树会被分割成若干个不相交的连通块（子树）。

假设某个连通块的大小为 $S$，那么完全在这个连通块内部的路径都不会经过颜色为 $c$ 的节点，这样的路径数量为 $\frac{S(S-1)}{2}$。

因此，颜色 $c$ 的贡献为：

$$\text{Contribution}(c) = \frac{n(n-1)}{2} - \sum_{i} \frac{S_i(S_i-1)}{2}$$

其中 $S_i$ 是删除颜色 $c$ 的节点后，各个连通块的大小。

可以使用树形 $dp$ 计算连通块的大小。

设 $f[c]$，表示目前遍历过程中，颜色为 $c$ 的节点已经分隔的连通块节点总数。

当我们遍历到一个节点 $u$（颜色为 $c_u$），并准备进入它的子树 $v$ 时：

• 记录下进入 $v$ 之前的分隔数：$pre=f[c_u]$。
• 递归遍历子树 $v$。
• 遍历完 $v$ 后，子树 $v$ 的总大小为 $siz[v]$。在子树 $v$ 内部，有一些节点已经被子树内部的 $c_u$ 颜色节点分隔了，分隔的数量就是 $f[c_u]-pre$。
• 那么，子树 $v$ 中没有被分隔且能一路连通到节点 $u$ 的节点数量就是 $cntv=siz[v] - (f[c_u] - pre)$。这些节点就构成了一个不包含颜色 $c_u$ 的连通块。
• 我们从总路径中减去这个连通块内部的路径数 $\frac{\text{cnt}\cdot(\text{cnt}-1)}{2}$，并将这部分大小累加到 $f[c_u]$ 中。

当整棵树遍历完成后，每种颜色可能在靠近树根的地方还剩下一个“顶部连通块”。对于颜色 $c$，这个顶部连通块的大小为 $n - f[c]$，同样需要将其内部路径数减去。