【启发式合并】（dsu on tree）讲解+例题
超级好的讲解（来自cf）
一、启发式合并的作用：
我们可以在O(nlogn) 的时间内，回答下列形式的所有询问：
节点v的子树中有多少节点有某种特性。
例如：
给一棵树，每个节点有一种颜色，询问：节点v的子树中，有多少节点的颜色是color c?
二、解决方法
2.1 准备工作：计算每个节点的子树大小
int sz[maxn];
void getsz(int v, int p)
{
sz[v] = 1; // every vertex has itself in its subtree
for(auto u : g[v])
if(u != p)
{
getsz(u, v);
sz[v] += sz[u]; // add size of child u to its parent(v)
}
}
2.2 普通解法（O(N ^ 2) ）
int cnt[maxn];
void add(int v, int p, int x)
{
cnt[ col[v] ] += x;
for(auto u: g[v])
if(u != p)
add(u, v, x)
}
void dfs(int v, int p)
{
add(v, p, 1);
//now cnt[c] is the number of vertices in subtree of vertex v that has color c. You can answer the queries easily.
add(v, p, -1);
for(auto u : g[v])
if(u != p)
dfs(u, v);
}
个人理解：要求得每一棵子树的值，那么就需要对每一棵子树进行遍历，由于只有一个cnt记录空间，因此每次遍历完后，需要将cnt数组清空，然后才能进行下一个兄弟节点的遍历，否则会将从cnt里的数据搞混淆，得出错误的结果。
这样一个节点会被重复遍历多次，复杂度很高。
2.3 easy to code but O(n log^2 n)
void dfs(int v, int p)
{
int mx = -1, bigChild = -1;
for(auto u : g[v])
if(u != p)
{
dfs(u, v);
if(sz[u] > mx)
mx = sz[u], bigChild = u;
}
if(bigChild != -1)
cnt[v] = cnt[bigChild];
else
cnt[v] = new map<int, int> ();
(*cnt[v])[ col[v] ] ++;
for(auto u : g[v])
if(u != p && u != bigChild)
{
for(auto x : *cnt[u])
(*cnt[v])[x.first] += x.second;
}
//now (*cnt[v])[c] is the number of vertices in subtree of vertex v that has color c. You can answer the queries easily.
}
2.4 easy to code and O(nlogn)
vector<int> *vec[maxn];
int cnt[maxn];
void dfs(int v, int p, bool keep)
{
int mx = -1, bigChild = -1;
for(auto u : g[v])
if(u != p && sz[u] > mx)
mx = sz[u], bigChild = u;
for(auto u : g[v])
if(u != p && u != bigChild)
dfs(u, v, 0);
if(bigChild != -1)
dfs(bigChild, v, 1), vec[v] = vec[bigChild];
else
vec[v] = new vector<int> ();
vec[v]->push_back(v);
cnt[ col[v] ]++;
for(auto u : g[v])
if(u != p && u != bigChild)
for(auto x : *vec[u])
{
cnt[ col[x] ]++;
vec[v] -> push_back(x);
}
//now (*cnt[v])[c] is the number of vertices in subtree of vertex v that has color c. You can answer the queries easily.
// note that in this step *vec[v] contains all of the subtree of vertex v.
if(keep == 0)
for(auto u : *vec[v])
cnt[ col[u] ]--;
}
2.5 heavy-light decomposition style O(nlogn) 【重点】
int cnt[maxn];
bool big[maxn];
void add(int v, int p, int x)
{
cnt[ col[v] ] += x;
for(auto u: g[v])
if(u != p && !big[u])
add(u, v, x)
}
void dfs(int v, int p, bool keep)
{
int mx = -1, bigChild = -1;
for(auto u : g[v])
if(u != p && sz[u] > mx)
mx = sz[u], bigChild = u;
for(auto u : g[v])
if(u != p && u != bigChild)
dfs(u, v, 0); // run a dfs on small childs and clear them from cnt
if(bigChild != -1)
dfs(bigChild, v, 1), big[bigChild] = 1; // bigChild marked as big and not cleared from cnt
add(v, p, 1);
//now cnt[c] is the number of vertices in subtree of vertex v that has color c. You can answer the queries easily.
if(bigChild != -1)
big[bigChild] = 0;
if(keep == 0)
add(v, p, -1);
}
个人理解：
- （原始解法）要求得每一棵子树的值，那么就需要对每一棵子树进行遍历，由于只有一个cnt记录空间，因此每次遍历完后，需要将cnt数组清空，然后才能进行下一个兄弟节点的遍历，否则会将从cnt里的数据搞混淆，得出错误的结果。
- 我们发现其实这里面有些重复的步骤可以省略。比如说，在遍历一个节点u时，因为是dfs，所以一定是先将u的所有子节点都dfs一遍后，然后才会返回到u，对u进行求解。而对u进行求解时，要对子节点求和再加上他自己，还是需要遍历它的所有子节点，那么既然上面已经遍历一遍子节点，那么刚刚遍历的子节点信息不删除，直接在这里用，不就节省了一遍重复遍历的过程吗？
- 但是不能保留下所有的子节点信息。因为我们进行dfs时，是公用cnt这个数组的，如果不将之前的信息清除掉，就会造成混乱。（比如节点1有三个子节点2、3、4，而4有两个子节点5、6。现在对1进行遍历，那么首先要对子节点遍历，也就是进行后序遍历：2、3、5、6、4、1。如果对2进行遍历后，不清空cnt，那么接着对3进行遍历的时候，就会把2和3的信息混合，导致3的信息出错。）
- 但是我们可以保留最后一个遍历的子节点的信息，然后再把前面清空掉的信息，再遍历一遍，这样相当于少遍历了一遍。（还是上面那个例子，我们遍历完4的信息后，马上就要返回1了，而计算1的时候，还是需要对2、3、4计算一遍，因为1 = 1 + 2 + 3+ 4。那么从4返回1的时候，4的信息就不用清除了，此时cnt中装有4的信息，即cnt[1] = cnt[4]。然后再把2、3和1自己的信息加上，就完成了对1的遍历。这个过程中可以看到，和原始做饭做法相比，少计算了一次4的值）。
- 既然可以少算一个子节点的值，那么选择哪个节点作为这个最后一个即可以被少算的这个节点，是至关重要的。显然，我们应该选择规模最大的一个子节点（也叫作重儿子），这样就可以少花点时间。
- 于是可以得出算法：每一次先对轻儿子进行遍历，求出轻儿子的ans，在这个过程中，每一次都要清空cnt。然后计算重儿子的ans，计算完后，不对cnt清空，接着求父节点自己的，再把轻儿子的值再遍历一遍，累加到cnt中。这样就完成了对一整棵树所有节点ans的遍历求解。
- 和原始做法相比，我们节约了对重儿子的重复计算。
自己修改过的板子
//CodeForces-600E
#include <iostream>
#include <stdio.h>
#include <string.h>
#include <string>
#include <math.h>
#include <algorithm>
#include <vector>
using namespace std;
const int maxn = 1e5 + 100;
int col[maxn], sz[maxn], son[maxn], cnt[maxn];
int n, big, maxx;
long long ans[maxn], sum;
vector<int> G[maxn];
void get_son(int u, int p)
{
sz[u] = 1;
int len = G[u].size();
for(int i = 0; i < len; i++)
{
int v = G[u][i];
if(v != p)
{
get_son(v, u);
sz[u] += sz[v];
if(sz[v] > sz[son[u]])
son[u] = v;
}
}
}
void add(int u, int p, int x)
{
cnt[ col[u] ] += x;
if(cnt[ col[u] ] > maxx)
{
sum = col[u];
maxx = cnt[ col[u] ];
}
else if(cnt[ col[u] ] == maxx)
{
sum += col[u];
}
int len = G[u].size();
for(int i = 0; i < len; i++)
{
int v = G[u][i];
if(v != p && v != big)
add(v, u, x);
}
}
void dfs(int u, int p, bool keep)
{
int len = G[u].size();
for(int i = 0; i < len; i++)
{
int v = G[u][i];
if(v != p && v != son[u])
dfs(v, u, 0);
}
if(son[u])
{
dfs(son[u], u, 1);
big = son[u];
}
add(u, p, 1);
big = 0;
ans[u] = sum;
if(!keep)
{
add(u, p, -1);
maxx = 0;
sum = 0;
}
}
int main()
{
while(~scanf("%d", &n))
{
big = 0;
maxx = 0;
sum = 0;
for(int i = 1; i <= n; i++)
{
G[i].clear();
sz[i] = son[i] = cnt[i] = ans[i] = 0;
}
for(int i = 1; i <= n; i++)
scanf("%d", &col[i]);
int u, v;
for(int i = 0; i < n - 1; i++)
{
scanf("%d%d", &u, &v);
G[u].push_back(v);
G[v].push_back(u);
}
get_son(1, 0);
dfs(1, 0, 1);
for(int i = 1; i <= n; i++)
{
if(i != 1)
printf(" ");
printf("%lld", ans[i]);
}
printf("\n");
}
return 0;
}