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完整代码如下：
代码语言：javascript复制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)
}
}
}
