树上启发式合并(dsu on tree)

前置概念阐述：
启发式：启发式算法是基于人类的经验和直观感觉，对一些算法的优化
启发式合并:对于两个大小不一样的集合，我们将小的集合合并到大的集合中，而不是将大的集合合并到小的集合中。这是基于人类经验和推理得出的原有合并算法的优化
轻重链、轻儿子、重儿子等概念详见树链剖分树链剖分oi-wiki

---

针对某些树上离线查询问题
树上启发式合并思路：

1. 第一次dfs划分轻重链，轻儿子以及重儿子
2. 第二次dfs先对当前节点轻儿子对应的子树进行处理，得到该子树的每个点的答案；同时消除以轻儿子为根的子树对原集合统计量的影响；之后再处理重儿子对应的子树所有节点答案，但不消除对其子树进行统计后带给原集合的影响。最后再次统计所有轻边对应的子树得到当前节点的答案。
   复杂度分析：
   已知在节点数为n的树上从根节点到任意节点路径上的轻边个数为$O(logn)$级别，根据我们进行启发式合并的思路，从任意一个节点到根节点的路径上有多少条轻边，就会被遍历多少次，也就是说每个节点被遍历的次数为$O(logn)$级别，共n个节点，则总时间复杂度为$O(nlogn)$级别。也即我们的深度优先遍历递归次数为$O(nlogn)$级别。（在第二次dfs中，一般来说除了递归操作，其中只允许$O(1)$的操作）

问题1：树上数颜色

代码：

#include<bits/stdc++.h>
#define maxn 100500
#define INF 0x3f3f3f3f
#define mem(a,b) memset(a,b,sizeof(a))
#define ll long long
using namespace std;
struct Edge{
    int to,next;
}edge[maxn<<2];
int c[maxn],head[maxn],n;
int sz[maxn],son[maxn];
int cnt[maxn],clr=0,ans[maxn];
int tot=0;
void addedge(int u,int v)
{
    edge[++tot].to=v;
    edge[tot].next=head[u];
    head[u]=tot;
} 
void dfs(int p,int fa,int fl)
{
    cnt[c[p]]+=fl;
    if(cnt[c[p]]==0&&fl==-1)clr--;
    if(cnt[c[p]]==1&&fl==1)clr++;
    for(int i=head[p];~i;i=edge[i].next)
    {
        int to=edge[i].to;
        if(to==fa)continue;
        dfs(to,p,fl);
    }
}
void dfs1(int p,int fa)
{
    sz[p]=1,son[p]=0;
    for(int i=head[p];~i;i=edge[i].next)
    {
        int to=edge[i].to;
        if(to==fa)continue;
        dfs1(to,p),sz[p]+=sz[to];
        if(sz[son[p]]<sz[to])son[p]=to;
    }
}
void dfs2(int p,int fa)
{
    for(int i=head[p];~i;i=edge[i].next)
    {
        int to=edge[i].to;
        if(to==fa||to==son[p])continue;
        dfs2(to,p);
        dfs(to,p,-1);//进行轻儿子对cnt数组影响的消除
    }
    if(son[p])dfs2(son[p],p);
    cnt[c[p]]++;
    if(cnt[c[p]]==1)clr++;
    for(int i=head[p];~i;i=edge[i].next)
    {
        int to=edge[i].to;
        if(to==fa||to==son[p])continue;
        dfs(to,p,1);//进行cnt数组的更新
    }
    ans[p]=clr;
}
int main()
{   
    scanf("%d",&n);
    mem(sz,0),mem(son,0);
    mem(cnt,0);
    for(int i=0;i<=n;i++)head[i]=-1;
    for(int i=1;i<=n-1;i++)
    {
        int u,v;
        scanf("%d%d",&u,&v);
        addedge(u,v);
        addedge(v,u);
    }
    for(int i=1;i<=n;i++)scanf("%d",&c[i]);
    dfs1(1,1);
    dfs2(1,1);
    int m;
    scanf("%d",&m);
    for(int i=1;i<=m;i++)
    {
        int x;
        scanf("%d",&x);
        printf("%d\n",ans[x]);
    }
    system("pause");
}

问题2：颜色平衡树

代码：

#include<bits/stdc++.h>
#define maxn 200500
#define INF 0x3f3f3f3f
#define mem(a,b) memset(a,b,sizeof(a))
#define ll long long
using namespace std;
int c[maxn],n,head[maxn],sz[maxn],son[maxn],tot=0;
struct Edge{
    int to,next;
}edge[maxn];
int mi=INF,mx=0,cnt[2][maxn],ans[maxn];
//cnt[0][x]：表示元素 x 的当前频率。
//cnt[1][f]：表示频率为 f 的元素个数。
void addedge(int u,int v)
{
    edge[++tot].to=v;
    edge[tot].next=head[u];
    head[u]=tot;
}
void add(int cl)
{
    cnt[1][cnt[0][cl]]--;
    cnt[0][cl]++;
    cnt[1][cnt[0][cl]]++;
    if(cnt[0][cl]<mi)mi=cnt[0][cl];
    if(cnt[0][cl]>mx)mx=cnt[0][cl];
    if(!cnt[1][cnt[0][cl]-1])mi++;
}
void del(int cl)
{
    cnt[1][cnt[0][cl]]--;
    cnt[0][cl]--;
    cnt[1][cnt[0][cl]]++;
    if(cnt[0][cl]<mi&&cnt[0][cl])mi=cnt[0][cl];
    if(cnt[0][cl]>mx)mx=cnt[0][cl];
    if(!cnt[1][cnt[0][cl]+1])mx--;
}
void dfs(int p,int fa,int fl)
{
    if(fl==-1)del(c[p]);
    else add(c[p]);
    for(int i=head[p];~i;i=edge[i].next)
    {
        int to=edge[i].to;
        if(to==fa)continue;
        dfs(to,p,fl);
    }
}
void dfs1(int p,int fa)
{
    sz[p]=1,son[p]=0;
    for(int i=head[p];~i;i=edge[i].next)
    {
        int to=edge[i].to;
        if(to==fa)continue;
        dfs1(to,p),sz[p]+=sz[to];
        if(sz[son[p]]<sz[to])son[p]=to;
    }
}
void dfs2(int p,int fa)
{
    for(int i=head[p];~i;i=edge[i].next)
    {
        int to=edge[i].to;
        if(to==fa||to==son[p])continue;
        dfs2(to,p);
        dfs(to,p,-1);
    }
    if(son[p])dfs2(son[p],p);
    add(c[p]);
    for(int i=head[p];~i;i=edge[i].next)
    {
        int to=edge[i].to;
        if(to==fa||to==son[p])continue;
        dfs(to,p,1);
    }
    ans[p]=(mi==mx);
}
int main()
{
    scanf("%d",&n);
    mem(sz,0),mem(son,0);
    mem(cnt,0),mem(ans,0);
    for(int i=0;i<=n;i++)head[i]=-1;
    for(int i=1;i<=n;i++)
    {
        int u;
        scanf("%d",&c[i]);
        scanf("%d",&u);
        if(u==0)u=1;
        addedge(u,i);
    }
    dfs1(1,1);
    dfs2(1,1);
    int tt=0;
    for(int i=1;i<=n;i++)
    tt+=ans[i];
    printf("%d",tt);
    system("pause");
}