题意

给你一棵树(即一个没有循环的连接的无向图),由编号从0到n-1的n个节点和正好n-1条边组成。树的根是节点0,树的每个节点都有一个标签,是字符串labels中给出的小写字符(即数字为i的节点有标签labels[i])。

边缘数组的形式是edges[i] = [ai, bi],这意味着树上的节点ai和bi之间有一条边缘。

返回一个大小为n的数组,其中ans[i]是第i个节点的子树中与节点i有相同标签的节点数。

一棵树T的子树是由T中的一个节点及其所有的子节点组成的树。

思路

这道题就是dfs自底向上,对于每个节点,分别统计他的子树中每个字符的数量。然后道当前节点后,将数量加到自身,然后返回。

labels = abaedcd

Untitled

按照思路以右侧为例,从底向上计算

首先节点3,定义一个26个元素的数组,因为没有任何子树,所以与他有相同标签的只有他自己。返回 一个数组,标记e有一个。

同理节点6,定义一个26个元素的数组,因为没有任何子树,所以与他有相同标签的只有他自己。返回一个数组,标记d有一个。

对于2节点,定义一个26个元素的数组total

先得到节点3的返回值,然后的得到节点6的返回值,将这两个数组中的标签数量分别加到total,

得到a: 1, d:1, e:1 此时我们就可以知道与节点2标签相同的节点数量了。将数组返回,再由节点1继续同样的计算,最后得到a: 2, d:1, e;1.

func countSubTrees(n int, edges [][]int, labels string) []int {
	ans := make([]int, n)
	graph := make(map[int][]int)
	for _, e := range edges {
		if _, ok := graph[e[0]]; !ok {
			graph[e[0]] = make([]int, 0)
		}
		if _, ok := graph[e[1]]; !ok {
			graph[e[1]] = make([]int, 0)
		}
		graph[e[0]] = append(graph[e[0]], e[1])
		graph[e[1]] = append(graph[e[1]], e[0])
	}

	var dfs func(int, int) []int
	dfs = func(node, parent int) []int {
		rels := graph[node]
		total := make([]int, 26)
		total[labels[node]-'a']++
		for _, next := range rels {
			if next == parent {
				continue
			}
			next := dfs(next, node)
			for i, n := range next {
				if n != 0 {
					total[i] += n
				}
			}
		}
		ans[node] = total[labels[node]-'a']
		return total
	}
	dfs(0, -1)
	return ans
}