Daily LeetCode Challenge - 1519. Number of Nodes in the Sub-Tree With the Same Label

Problem From.
https://leetcode.com/problems/number-of-nodes-in-the-sub-tree-with-the-same-label/

오늘 문제는 edges 배열이 주어지면 tree 를 만들어서 각 node 의 child 노드에 parent 노드와 같은 알파벳이 몇개 있는지 세어서 intarray 를 반환하는 문제였다.

기본적으로 dfs 를 쓰고, 26개짜리 알파벳 길이 만큼의 labesCount 배열을 각 노드를 방문할때 마다 활용하여, 그 labesCount 배열 안에 각 알파벳이 지금까지 몇개가 나왔는지를 누적해주는게 핵심이었다.

class Solution {
    lateinit var answer : IntArray
    val edgesMap = HashMap<Int, MutableList<Int>>()
    lateinit var labelMap : CharArray
    lateinit var visited : HashSet<Int>
    
    fun countSubTrees(n: Int, edges: Array<IntArray>, labels: String): IntArray {
        answer = IntArray(n)
        labelMap = CharArray(n)
        visited = HashSet<Int>(n)
        
        
        edges.forEach {
            edgesMap.putIfAbsent(it[0], mutableListOf<Int>())
            edgesMap.putIfAbsent(it[1], mutableListOf<Int>())            
            edgesMap.get(it[0])?.add(it[1])
            edgesMap.get(it[1])?.add(it[0])
        }
        
        labels.forEachIndexed { i, char -> 
            labelMap[i] = char
        }
        dfs(0)
        return answer
    }
    
    fun dfs(node: Int): IntArray {
        var labelsCount: IntArray = IntArray(26)
        
        visited.add(node)
        for (child in edgesMap.getOrDefault(node, mutableListOf<Int>())) {
            if(child in visited) continue
            
            val labels = dfs(child)
            
            for (i in 0..25) {
                labelsCount[i] += labels[i]
            }
        }
        labelsCount[labelMap[node] - 'a']++
        
        answer[node] = labelsCount[labelMap[node] - 'a']
        return labelsCount
        
    }
}