Daily LeetCode Challenge - 863. All Nodes Distance K in Binary Tree

Problem From.

https://leetcode.com/problems/all-nodes-distance-k-in-binary-tree/

오늘 문제는 target node 로 부터 주어진 거리 내에 있는 모든 원소를 찾아서 반환하는 문제였다.

이 문제는 BFS 로 풀 수 있었는데, 먼저 서로 tree 를 graph 형식으로 바꾸어서 서로 연결되어 있는 상황을 나타내준다. 그런 다음 각각의 연결된 관계를 따라서 target node 를 찾은 다음 BFS 를 통해서 인접한 노드를 arrayList 에 넣어서 반환해주면 되었다.

/**
 * Definition for a binary tree node.
 * class TreeNode(var `val`: Int = 0) {
 *     var left: TreeNode? = null
 *     var right: TreeNode? = null
 * }
 */

import java.util.*

class Solution {
    
    val map = HashMap<TreeNode, ArrayList<TreeNode>>()
    
    fun distanceK(root: TreeNode?, target: TreeNode?, K: Int): List<Int> {


        makeGraph(root, null, target)

        val queue = LinkedList<TreeNode>()
        val visited = HashSet<TreeNode>()
        queue.offer(target)
        visited.add(target!!)

        var distance = 0
        while (!queue.isEmpty()) {
            val size = queue.size
            val answer = ArrayList<Int>()
            for (i in 0..size - 1) {
                val n = queue.poll()
                answer.add(n.`val`)
                if (map.contains(n)) {
                    for (num in map[n]!!) {
                        if (!visited.contains(num)) {
                            queue.offer(num)
                            visited.add(num)
                        }
                    }
                }
            }
            if (distance++ == K) return answer
        }

        return listOf()
    }
    
    fun makeGraph(root: TreeNode?, parent: TreeNode?, target: TreeNode?) {
        if (root == null) return
        if (parent != null) {
            map.getOrPut(root, { arrayListOf() }).add(parent)
            map.getOrPut(parent, { arrayListOf() }).add(root)
        }
        makeGraph(root.left, root, target)
        makeGraph(root.right, root, target)
    }
}