2024-04-21：用go语言，给一棵根为1的树，每次询问子树颜色种类

2024-04-21：用go语言，给一棵根为1的树，每次询问子树颜色种类数。

假设节点总数为n，颜色总数为m，

每个节点的颜色，依次给出，整棵树以1节点做头，

有k次查询，询问某个节点为头的子树，一共有多少种颜色。

1 <= n, m, k <= 10^5。

答案2024-04-21：

来自左程云。

chatgpt

大体步骤如下：

大体过程描述：

1.数据结构初始化：定义全局变量和数组用来存储图的结构、节点颜色等信息，并初始化相关数组和变量。

2.输入处理：通过预定义的输入数组，按给定格式依次读取节点数n，建立树的连接关系，记录每个节点的颜色。

3.DFS遍历：

• 第一次DFS（dfs1）：计算每个节点子树的大小，并标记每个节点的重节点。
• 第二次DFS（dfs2）：处理每个节点的子树，包括处理重节点和非重节点的不同子树，更新颜色计数和子树的颜色种类数。

4.颜色计数：通过add函数和delete函数实现颜色的增加与减少操作，维护当前节点子树中颜色种类的计数。

5.输出查询结果：对于每次查询，按照给定节点进行处理，并输出计算得到的颜色种类数。

时间复杂度：

• DFS1：对整个树进行一次DFS，时间复杂度为O(n)。
• DFS2：同样对整个树进行一次DFS，时间复杂度为O(n)。
• add和delete函数：每个节点至多被遍历4次（每条边两次），因此每次add和delete的时间复杂度为O(n)。
• 查询：对于每次查询，计算颜色种类数时需要遍历整个子树，时间复杂度为O(n)。

综上，总的时间复杂度为O(n)。

空间复杂度：

• graph, color, size, heavy, cnt, ans：每个数组的长度为n，因此空间复杂度为O(n)。
• 其他局部变量：不超过常数大小，可忽略。

综上，总的额外空间复杂度为O(n)。

Go完整代码如下：

packagemain
import(
"fmt"
)
varMAXNint=200005
vargraph[][]int
varcolor[]int
varsize[]int
varheavy[]int
varcnt[]int
varans[]int
varn,mint
funcmain(){
graph=make([][]int,MAXN)
fori:=rangegraph{
graph[i]=make([]int,0)
}
color=make([]int,MAXN)
size=make([]int,MAXN)
heavy=make([]int,MAXN)
cnt=make([]int,MAXN)
ans=make([]int,MAXN)
inputs:=[]int{5,
1,2,
1,3,
2,4,
2,5,
1,2,2,3,3,
5,
1,
2,
3,
4,
5}
ii:=0
foriin=inputs[ii]
ii++
fori:=1;i<=n;i++{
graph[i]=make([]int,0)
}
fori:=1;ia:=inputs[ii]
ii++
b:=inputs[ii]
ii++
graph[a]=append(graph[a],b)
graph[b]=append(graph[b],a)
}
fori:=1;i<=n;i++{
c:=inputs[ii]
ii++
color[i]=c
}
dfs1(1,0)
dfs2(1,0,false)
m=inputs[ii]
ii++
fori:=1;i<=m;i++{
q:=inputs[ii]
ii++
fmt.Println(ans[q])
}
}
}
vartotalint
funcdfs1(cur,fatherint){
size[cur]=1
for_,next:=rangegraph[cur]{
ifnext!=father{
dfs1(next,cur)
size[cur]+=size[next]
ifsize[heavy[cur]]heavy[cur]=next
}
}
}
}
funcdfs2(cur,fatherint,isHeavybool){
for_,next:=rangegraph[cur]{
ifnext!=father&&next!=heavy[cur]{
dfs2(next,cur,false)
}
}
ifheavy[cur]!=0{
dfs2(heavy[cur],cur,true)
}
add(cur,father,heavy[cur])
ans[cur]=total
if!isHeavy{
delete(cur,father)
total=0
}
}
funcadd(cur,father,exceptint){
cnt[color[cur]]++
ifcnt[color[cur]]==1{
total++
}
for_,next:=rangegraph[cur]{
ifnext!=father&&next!=except{
add(next,cur,except)
}
}
}
funcdelete(cur,fatherint){
cnt[color[cur]]--
for_,next:=rangegraph[cur]{
ifnext!=father{
delete(next,cur)
}
}
}