【Luogu U41492】树上数颜色——树上启发式合并（dsu on tree）

（这题在洛谷主站居然搜不到……还是在百度上偶然看到的）

题目描述

给一棵根为1的树，每次询问子树颜色种类数

输入输出格式

输入格式：

第一行一个整数n，表示树的结点数

接下来n-1行，每行一条边

接下来一行n个数，表示每个结点的颜色c[i]

接下来一个数m，表示询问数

接下来m行表示询问的子树

输出格式：

对于每个询问，输出该子树颜色数

说明

对于前三组数据，1<=m,c[i]<=n<=100

1<=m,c[i]<=n<=1e5

　　本来在学树上启发式合并，偶然看到这个题，就当模板来打了。

　　树上启发式合并（dsu on tree）是一种对静态子树节点信息的统计手段。和树上点分治一样，二者都是经过优化的暴力做法，启发式合并的复杂度可以从平方级别降到O(nlogn)。

　　启发式合并的思想源自暴力统计子树信息的过程。这题中，某棵树的每个点都有一个颜色，我们要多组询问统计某个点为根的子树内有多少种颜色。

　　这不是主席树乱搞嘛

　　考虑直接做这个题的话，我们想预处理出每个点的答案，可以直接从每个点出发做一次dfs统计得出答案后再dfs一遍，把它清空。这个显然的做法是O(n^2)的。

　　dsu on tree给出的优化策略是，我们想要统计某个点u，可以先遍历完它的轻儿子们，把这些点清空消去对重儿子的影响，后再去统计它的重儿子。这样得到的重儿子的信息我们还可以利用在u上，所以这时候我们保留重儿子的信息，再暴力地把它的轻儿子统计一遍来更新u的答案就可以了。

　　直观上看，每次特殊对待（只统计一次）的这个节点选择重儿子显然是最优的。启发式合并的复杂度经过证明可以达到O(nlogn)；考虑每个点，这个点会被统计多少次呢？

　　1、通过轻边被扫到。根据轻重链的性质，这一步最多会有log次。

　　2、被遍历被扫到一次，或者藉由所在重链被扫到一次。是的，我们在合并的过程中遍历重链就像树剖的第二次dfs一样，每条重链只会被扫一次，因此这两种情况是O(1)的。

　　dsu on tree的复杂度得到了证明，我们可以愉快地用它来套做题了。需要强调的是，这个算法适用的是“静态统计子树信息”的问题。具体的遍历写法见代码：

1. #include <iostream>
2. #include <cstdio>
3. #include <cstring>
4. #include <algorithm>
5. #include <cctype>
6. #define maxn 100010
7. using namespace std;
8. template <class T>
9. void read(T &x) {
10. x = 0;
11. char ch = getchar();
12. while (!isdigit(ch))
13. ch = getchar();
14. while (isdigit(ch)) {
15. x = x * 10 + (ch ^ 48);
16. ch = getchar();
17. }
18. }
19. int head[maxn], top = 1;
20. struct E {
21. int to, nxt;
22. } edge[maxn << 1];
23. inline void insert(int u, int v) {
24. edge[++top] = (E) {v, head[u]};
25. head[u] = top;
26. }
27. int n, m;
28. int cnt[maxn], color[maxn], son[maxn], size[maxn], ans[maxn], NowSon, sum;
29. void dfs1(int u, int pre) {  //预处理子树大小
30. size[u] = 1;
31. for (int i = head[u]; i; i = edge[i].nxt) {
32. int v = edge[i].to;
33. if (v == pre) continue;
34. dfs1(v, u);
35. size[u] += size[v];
36. if (size[son[u]] < size[v])
37. son[u] = v;
38. }
39. }
40. void cal(int u, int pre, int val) {  //计算答案
41. if (!cnt[color[u]]) ++sum;//在当前树中统计这种颜色
42. cnt[color[u]] += val;
43. for (int i = head[u]; i; i = edge[i].nxt) {
44. int v = edge[i].to;
45. if (v == pre || v == NowSon)//避开根的重儿子
46. continue;
47. cal(v, u, val);
48. }
49. }
50. void dsu(int u, int pre, bool op) {//启发式合并的主体。op为 0表示这次操作由轻边遍历得到，需要清空
51. for (int i = head[u]; i; i = edge[i].nxt) {
52. int v = edge[i].to;
53. if (v == pre || v == son[u])
54. continue;
55. dsu(v, u, 0);//轻儿子
56. }
57. if (son[u])
58. dsu(son[u], u, 1), NowSon = son[u];//重儿子
59. cal(u, pre, 1); NowSon = 0;//重新扫描轻儿子，二次统计
60. ans[u] = sum;//那么现在的颜色数就是u的信息
61. if (!op) {//清空当前的统计数
62. cal(u, pre, -1);
63. sum = 0;
64. }
65. }
66. int main() {
67. read(n);
68. int u, v;
69. for (int i = 1; i < n; ++i) {
70. read(u), read(v);
71. insert(u, v), insert(v, u);
72. }
73. for (int i = 1; i <= n; ++i)
74. read(color[i]);
75. dfs1(1, 0);
76. dsu(1, 0, 1);
77. read(m);
78. for (int i = 1; i <= m; ++i) {
79. read(v);
80. printf("%d\n", ans[v]);
81. }
82. return 0;
83. }