题解：P9233 [蓝桥杯 2023 省 A] 颜色平衡树

本题要求统计一棵根树中有多少棵子树满足“出现过的每种颜色次数全部相等”。朴素地为每个结点开哈希表统计其整棵子树颜色分布复杂度会达到 \(O(n^2)\)，显然无法通过规模上限 \(n\le 2\times10^{5}\)。关键在于如何把子树统计做成一次性、且只在必要位置增减。

先对整棵树做一次 DFS 获得 Euler 序，用数组 \(\text{od}[1\ldots n]\) 记录节点在序中的出现顺序，并得到每个结点的子树区间 \([st[u],\,ed[u]]\)。随后用 DSU on Tree 自顶向下遍历：处理完一条重链后，把轻儿子的小计全部并入这条重链对应的桶里，保证同一层级的插入总在“重”集合上完成，从而做到每个结点的贡献只在其所在重链被完整保留一次，其他场合被马上清空，整体复杂度 \(O(n)\)。

计数时维护两张整型全局数组：\(\text{cnt}[c]\) 表示当前处理的子树内颜色 \(c\) 的出现次数，\(\text{frq}[k]\) 表示“出现恰好 \(k\) 次的颜色种数”。再用整数 \(\text{dif}\) 记录当前不同出现次数的种类数，即满足 \(\text{frq}[k]>0\) 的 \(k\) 的个数。插入或删除一个节点时，先把旧次数 \(o=\text{cnt}[c]\) 在 \(\text{frq}\) 中减一；若该次数消失则 \(\text{dif}\) 减一。然后更新新次数 \(o\pm1\)，若该次数此前为空则 \(\text{dif}\) 加一。这样每一步均为 \(O(1)\)。

长链遍历到某结点 \(u\) 时，它的整棵子树统计已就绪。若 \(\text{dif}=1\)，说明此时所有出现次数都相同，子树 \(u\) 即为颜色平衡树，答案自增。遍历结束后若当前子树不是重链保留下来的那一条，则将整个区间 \([st[u],\,ed[u]]\) 反向删除，恢复计数器，保证轻子树不会污染兄弟分支。

时间复杂度 \(O(n)\)。

link