1+ class Solution {
2+ 3+ fun dfs (
4+ current : Int ,
5+ labels : String ,
6+ neighbours : Array <List <Int >>,
7+ counts : Array <IntArray >,
8+ visited : BooleanArray
9+ ) {
10+ val count = counts[current]
11+ val label = labels[current]
12+ 13+ visited[current] = true
14+ neighbours[current].forEach { neighbour ->
15+ if (! visited[neighbour]) {
16+ dfs(neighbour, labels, neighbours, counts, visited)
17+ counts[neighbour].forEachIndexed { i, neighbourCount ->
18+ count[i] + = neighbourCount
19+ }
20+ }
21+ }
22+ 23+ count[label.toInt() - ' a' + = 1
24+ }
25+ 26+ fun countSubTrees (n : Int , edges : Array <IntArray >, labels : String ): IntArray {
27+ val result = IntArray (n) { 0 }
28+ val counts = Array <IntArray >(n) { IntArray (26 ) { 0 } }
29+ val neighbours = Array <List <Int >>(n) { mutableListOf<Int >() }
30+ val visited = BooleanArray (n) { false }
31+ 32+ edges.forEach { edge ->
33+ (neighbours[edge[0 ]] as MutableList <Int >).add(edge[1 ])
34+ (neighbours[edge[1 ]] as MutableList <Int >).add(edge[0 ])
35+ }
36+ 37+ dfs(0 , labels, neighbours, counts, visited)
38+ 39+ labels.forEachIndexed { i, label ->
40+ result[i] = counts[i][label.toInt() - ' a'
41+ }
42+ 43+ return result
44+ }
45+ 46+ }
0 commit comments