2024-04-21:用go语言,给一棵根为1的树,每次询问子树颜色种类数。 假设节点总数为n,颜色总数为m, 每个节点的颜色,依次给出,整棵树以1节点做头, 有k次查询,询问某个节点为头的子树,一共

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

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)
		}
	}
}