(复习)树上启发式合并(dsu on tree)入门U41492树上数颜色
主要思想是树的重轻儿子之分使得时间复杂度为o(nlogn),神奇
欲深入了解的这里:https://oi-wiki.org/graph/dsu-on-tree/
点击查看代码
#include <bits/stdc++.h>
using namespace std;
typedef struct edge//边结构体
{
int to,next;
}EDGE;
//边相关数组
EDGE e[100001<<1];
int head[100001];
int cnt;
//树大小及DFS序相关数组
int son[100001];//根节点为i的重儿子编号
int big[100001];//树的大小
int b[100001];//树节点的起始DFS序
int ed[100001];//树节点的最远DFS序
int match[100001];//DFS序对应的节点编号
int pos;//DFS序数
//颜色,答案相关数组
int c[100001];//节点颜色
int colornumber;//现在有的颜色数量
int times[100001];//每个颜色出现的次数
int ans[100001];//节点x的答案
void add(int x,int y)//建立边
{
e[++cnt].to = y;
e[cnt].next = head[x];
head[x] = cnt;
}
void add_color(int x)//加颜色
{
times[c[x]]++;
if(times[c[x]] == 1) colornumber++;
}
void del_color(int x)//删颜色
{
times[c[x]]--;
if(times[c[x]] == 0) colornumber--;
}
void DFS1(int x,int fa)//预处理DFS
{
big[x] = 1;//初始化树大小
match[++pos] = x;//DFS序匹配
b[x] = pos;//节点的DFS序
for (int i = head[x];i;i = e[i].next)//DFS遍历
{
if(e[i].to == fa) continue;
DFS1(e[i].to,x);
big[x] += big[e[i].to];//累加大小
if(son[x] == 0||big[e[i].to] > big[son[x]]) son[x] = e[i].to;//更新重儿子
}
ed[x] = pos;//记录最远DFS序
}
void DFS2(int x,int fa,bool keep)//树上启发式合并(dus on tree),keep为是否计算完要删除
{
for (int i = head[x];i;i = e[i].next)//先计算所有轻儿子的树答案
{
if(e[i].to == fa||e[i].to == son[x]) continue;//注意条件
DFS2(e[i].to,x,false);//继承删除
}
if(son[x]) DFS2(son[x],x,true);//有重儿子的话计算,并先不删除
for (int i = head[x];i;i = e[i].next)//因为之前的轻儿子已经删除,要重新计算
{
if(e[i].to == fa||e[i].to == son[x]) continue;
//由于儿子的DFS序连续,直接循环,若递归代码非常复杂!!!
for (int j = b[e[i].to];j<=ed[e[i].to];j++) add_color(match[j]);//注意是DFS序的配对节点
}
add_color(x);//再加上自己
ans[x] = colornumber;//记录答案
if(!keep)//如果要删除
{
for (int i = b[x];i<=ed[x];i++) del_color(match[i]);//从自己开始删
}
}
int main()
{
ios::sync_with_stdio(0);cin.tie(0);cout.tie(0);
int n,m,x,y;
cin>>n;
for (int i = 1;i<=n-1;i++)//建树
{
cin>>x>>y;
add(x,y);
add(y,x);
}
for (int i = 1;i<=n;i++) cin>>c[i];//输入颜色
DFS1(1,0);//预处理
DFS2(1,0,false);//dus on tree
cin>>m;
while(m--)//回答m次询问
{
cin>>x;
cout<<ans[x]<<"\n";//o(1)回答
}
return 0;
}
