-
Notifications
You must be signed in to change notification settings - Fork 0
/
38NodesInSubTreeLable.go
42 lines (38 loc) · 1.04 KB
/
38NodesInSubTreeLable.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
//https://leetcode.com/problems/number-of-nodes-in-the-sub-tree-with-the-same-label/description/
func countSubTrees(n int, edges [][]int, labels string) []int {
graph := make(map[int][]int)
buildGraph(edges, graph)
res := make([]int, n)
visit := make(map[int]bool)
buildRes(0, res, graph, labels, visit)
return res
}
func buildGraph(edges [][]int, graph map[int][]int) {
for _, edge := range edges {
a, b := edge[0], edge[1]
if _, ok := graph[a]; !ok {
graph[a] = []int{}
}
if _, ok := graph[b]; !ok {
graph[b] = []int{}
}
graph[a] = append(graph[a], b)
graph[b] = append(graph[b], a)
}
}
func buildRes(curr int, res []int, graph map[int][]int, labels string, visit map[int]bool) map[rune]int {
visit[curr] = true
occMap := make(map[rune]int)
for _, child := range graph[curr] {
if !visit[child] {
childMap := buildRes(child, res, graph, labels, visit)
//update occMap
for k, v := range childMap {
occMap[k] += v
}
}
}
occMap[rune(labels[curr])] += 1
res[curr] = occMap[rune(labels[curr])]
return occMap
}