2024-04-21：用go语言，给一棵根为1的树，每次询问子树颜色种类数。 假设节点总数为n，颜色总数为m， 每个节点的颜色，

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完整代码如下：

package main

import (

"fmt"

)

var MAXN int = 200005

var graph [][]int

var color []int

var size []int

var heavy []int

var cnt []int

var ans []int

var n, m int

func main() {

graph = make([][]int, MAXN)

for i := range graph {

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

for ii < len(inputs) {

n = inputs[ii]

ii++

for i := 1; i <= n; i++ {

graph[i] = make([]int, 0)

}

for i := 1; i < n; i++ {

a := inputs[ii]

ii++

b := inputs[ii]

ii++

graph[a] = append(graph[a], b)

graph[b] = append(graph[b], a)

}

for i := 1; i <= n; i++ {

c := inputs[ii]

ii++

color[i] = c

}

dfs1(1, 0)

dfs2(1, 0, false)

m = inputs[ii]

ii++

for i := 1; i <= m; i++ {

q := inputs[ii]

ii++

fmt.Println(ans[q])

}

}

}

var total int

func dfs1(cur, father int) {

size[cur] = 1

for _, next := range graph[cur] {

if next != father {

dfs1(next, cur)

size[cur] += size[next]

if size[heavy[cur]] < size[next] {

heavy[cur] = next

}

}

}

}

func dfs2(cur, father int, isHeavy bool) {

for _, next := range graph[cur] {

if next != father && next != heavy[cur] {

dfs2(next, cur, false)

}

}

if heavy[cur] != 0 {

dfs2(heavy[cur], cur, true)

}

add(cur, father, heavy[cur])

ans[cur] = total

if !isHeavy {

delete(cur, father)

total = 0

}

}

func add(cur, father, except int) {

cnt[color[cur]]++

if cnt[color[cur]] == 1 {

total++

}

for _, next := range graph[cur] {

if next != father && next != except {

add(next, cur, except)

}

}

}

func delete(cur, father int) {

cnt[color[cur]]--

for _, next := range graph[cur] {

if next != father {

delete(next, cur)

}

}

}

在这里插入图片描述