LeetCode Entry

863. All Nodes Distance K in Binary Tree

12.07.2023 medium 2023 kotlin

List of k distanced from target nodes in a Binary Tree

863. All Nodes Distance K in Binary Tree medium blog post substack image.png

Join me on Telegram

https://t.me/leetcode_daily_unstoppable/272

Problem TLDR

List of k distanced from target nodes in a Binary Tree

Intuition

There is a one-pass DFS solution, but it feels like too much of a corner cases and result handholding. A more robust way is to traverse with DFS and connect children nodes to parent, then send a wave from target at k steps.

Approach

Let’s build an undirected graph and do BFS.

  • don’t forget a visited HashSet

Complexity

  • Time complexity: \(O(n)\)

  • Space complexity: \(O(n)\)

Code


fun distanceK(root: TreeNode?, target: TreeNode?, k: Int): List<Int> {
    val fromTo = mutableMapOf<Int, MutableList<Int>>()
        fun dfs(node: TreeNode?, parent: TreeNode?) {
            node?.run {
                parent?.let {
                    fromTo.getOrPut(`val`) { mutableListOf() } += it.`val`
                    fromTo.getOrPut(it.`val`) { mutableListOf() } += `val`
                }
                dfs(left, this)
                dfs(right, this)
            }
        }
        dfs(root, null)
        return LinkedList<Int>().apply {
            val visited = HashSet<Int>()
                target?.run {
                    add(`val`)
                    visited.add(`val`)
                }
                repeat(k) {
                    repeat(size) {
                        fromTo.remove(poll())?.forEach { if (visited.add(it)) add(it) }
                    }
                }
            }
        }