hdu6035 Colorful Tree — 树形 DP / 补集计数

题意

给一棵 $n$ 个点的树，每个点有一种颜色 $c_{i}$ 。一条路径的权值 = 路径上不同颜色的种数。求所有 $(\binom{n}{2})$ 条路径的权值之和。多组数据，读到 EOF。

核心想法：补集（按颜色贡献）

对每一种颜色 $c$ ，统计 $c$ 出现在多少条路径里：

答案 = \sum_{c} (总路径数 - 不包含颜色 c 的路径数)

其中总路径数 $= (\binom{n}{2})$ 。设出现过的颜色种数为 $n u m$ ，则

答案 = n u m \cdot (\binom{n}{2}) - \sum_{c} (不包含颜色 c 的路径数)

把树上所有颜色为 $c$ 的点删掉，树会被拆成若干连通块。若这些连通块大小为 $s_{1}, s_{2}, \dots$ ，那么 不包含颜色 $c$ 的路径数 = $\sum_{i} (\binom{s_{i}}{2})$ （两端点落在同一个连通块内，就不会经过任何颜色 $c$ 的点）。

于是问题变成：对每种颜色，求出“删掉这种颜色的点后，各连通块的 $(\binom{s}{2})$ 之和”。朴素做法是枚举每种颜色重扫一遍， $O (n u m \cdot n)$ 会超时，需要一趟 DFS 同时完成所有颜色的统计。

一趟 DFS 的关键观察（“合并计数器”）

以 $1$ 为根做 DFS，维护两个数组：

sons[u]： $u$ 的子树大小；
sum[c]：对颜色 $c$ 来说，已经“合并/切掉”掉的节点数——也就是所有已经成块、且该块不是根所在连通块的节点数（含颜色 $c$ 本身的点）。

进入节点 $u$ （颜色 $c u$ ）时：

sons[u]=1，sum[c[u]]++（把 $u$ 归入颜色 $c [u]$ 的累计）；
记 pre = sum[c[u]] 作为处理 $u$ 的子节点前的基线；
对每个子节点 $v$ ，先递归处理完 $v$ 的子树，再：
- sons[u] += sons[v]；
- sum[c[u]] - pre 就是这棵子树里颜色 $c [u]$ 的点数（它们是“隔断点”）；
- 子树里不属于颜色 $c [u]$ 的、在最近的颜色 $c [u]$ 隔断点之上的点，构成一个不含根的连通块，大小为 ct = sons[v] - (sum[c[u]] - pre)：s += C(ct,2)，再 sum[c[u]] += ct；
- 更新 pre = sum[c[u]]，保证下一个子节点只统计它自己的部分，不重复前面兄弟子树。

这里 sum[c[u]] 会不断把切掉的块“吞并”进去，所以对根的颜色（比如整棵树的根点在 DFS 最后会把几乎所有节点并进去）， sum[c] 会一路增大。DFS 结束后：

s 已经累加了每种颜色下所有不含根的连通块的 $(\binom{s}{2})$ ；
对每种颜色 $c$ ，只剩一个含根连通块还没算，它的大小是 n - sum[c]，补上 C(n-sum[c], 2)。

因此最后答案：

text

ans = num * n*(n-1)/2 - s - Σ_{每种颜色 c} C(n - sum[c], 2)

一个小验证

样例 2（n=6, colors=1 2 1 3 2 1，根为 1，颜色 1）：

颜色 1：sum=6（根是颜色 1），n-sum=0，不贡献；
颜色 2：n-sum=3，C(3,2)=3；
颜色 3：n-sum=5，C(5,2)=10；
s=3（这是颜色 1 的那个不含根连通块 {2,4,5}）。

num*C(6,2) - s - (0+3+10) = 3*15 - 3 - 13 = 29，与样例一致。

复杂度

时间： $O (n)$ （一次遍历，每个点进出各一次）；
空间： $O (n)$ 。

实现细节 / 坑

多组数据，读入用 scanf 直到 EOF；每组要清空 sum、vis、邻接表。
$n \leq 2 \times 10^{5}$ ，答案与中间量都要用 long long。
递归 DFS 在一条链上会爆栈（实测 $n = 2 \times 10^{5}$ 链子上段错误），所以用显式栈 迭代模拟后序遍历 + 合并计数，既保证顺序与递归完全一致，又不会爆栈。

这题教会什么

补集计数：直接统计“包含颜色 $c$ 的路径”困难，就换成“总路径 − 不包含 $c$ 的路径”；
虚树思想：删掉同色点后按连通块计数，用一趟 DFS 的“合并计数器”把所有颜色的连通块同时算出来；
分治根块：把块分成“含根块”和“不含根块”两类，分别用 DFS 处理小部分、用 n - sum[c] 处理最后一个。