863.all-nodes-distance-k-in-binary-tree.md

863. All Nodes Distance K in Binary Tree

Test Cases

Normal Cases

Input: 
        1
      /   \
     2     3
   /   \ /   \
  4    5 6    7
        \
         8
target = 2, k = 2
Output: [8, 3]

Edge / Corner Cases

• k is 0.

Input: ...
Output: The target node itself. 

• Target node is a leaf node.

Input: 
        1
      /   \
     2     3
   /     /   \
  4      6    7
target = 2, k = 3
Output: [6, 7]

• There is no node at the distance k from the target node.

Input: 
        1
      /   \
     2     3
   /   \ /   \
  4    5 6    7
        \
         8
target = 2, k = 4
Output: []

Key Insights

1. We traverse the tree to update the parent node until we reach the target node.
2. Then we keep traversing the tree to find the nodes at the distance k from the target node.
3. We can use either DFS or BFS to traverse the tree.

• 由于是二叉链表，所以，无法直接由当前结点走向其父节点，所以用了一个map加一次dfs遍历来保存所有结点的父节点，这样就可以查表直接跳到父节点了。
• 为了防止走回头路，所以设计了一个from标志，等效于设置一个set（将路过的结点加入set，若待访问的结点不在set中，则访问它，否则跳过）。

BFS (Level by Level)

private val parentMap = HashMap<TreeNode, TreeNode>()

fun distanceK(root: TreeNode?, target: TreeNode?, k: Int): List<Int> {
    updateParent(root, target)
    return bfs(target, k)
}

private fun bfs(target: TreeNode?, k: Int): List<Int> {
    val ans = mutableListOf<Int>()
    if (target == null) return ans
    var distance = 0
    val queue = ArrayDeque<TreeNode>()
    val visited = HashSet<TreeNode>()
    queue.addLast(target)
    visited.add(target)
    while (queue.isNotEmpty() && distance <= k) {
        val size = queue.size
        repeat (size) {
            val node = queue.removeFirst()
            if (distance == k) {
                ans.add(node.`val`)
            }
            if (node.left != null && node.left !in visited) {
                queue.addLast(node.left)
                visited.add(node.left)
            }
            if (node.right != null && node.right !in visited) {
                queue.addLast(node.right)
                visited.add(node.right)
            }
            if (parentMap[node] != null) {
                val parent = parentMap[node]!!
                if (parent !in visited) {
                    queue.addLast(parent)
                    visited.add(parent)
                }
            }
        }
        distance++
    }
    return ans
}

// Or equivalently, we can enqueue the node and distance to the queue.
private fun bfs(target: TreeNode?, k: Int): List<Int> {
    val results = mutableListOf<Int>()
    if (target == null) return results
    val queue = ArrayDeque<Pair<TreeNode, Int>>()
    val visited = hashSetOf<Int>()

    queue.addLast(target to 0)
    visited.add(target.`val`)
    while (queue.isNotEmpty()) {
        val (node, distance) = queue.removeFirst()
        if (distance == k) {
            results.add(node.`val`)
        } else if (distance < k) {
            if (node.left != null && node.left !in visited) {
                queue.addLast(node.left to distance + 1)
                visited.add(node.left.`val`)
            }
            if (node.right != null && node.right !in visited) {
                queue.addLast(node.right to distance + 1)
                visited.add(node.right.`val`)
            }
            if (parentMap[node] != null) {
                val parent = parentMap[node]!!
                if (parent !in visited) {
                    queue.addLast(parent to distance + 1)
                    visited.add(parent.`val`)
                }
            }
        }
    }
    return results
}

• Time Complexity: O(n) where n is the number of nodes in the tree.
• Space Complexity: O(n) for the parent map and visited set.

DFS

private val results = mutableListOf<Int>()
private val parentMap = HashMap<TreeNode, TreeNode>()
private val visited = HashSet<TreeNode>()

fun distanceK(root: TreeNode?, target: TreeNode?, k: Int): List<Int> {
    if (root == null || target == null) return results
    updateParent(root, target)
    dfs(target, 0, k)
    return results
}

private fun updateParent(root: TreeNode?, target: TreeNode) {
    // We stop at the empty or target node
    if (root == null || root == target) return
    if (root.left != null) {
        parentMap[root.left] = root
        updateParent(root.left, target)
    }    
    if (root.right != null) {
        parentMap[root.right] = root
        updateParent(root.right, target)
    }
}

// Use visited set to prevent duplicate visiting.
private fun dfs(root: TreeNode?, distance: Int, k: Int) {
    if (root == null || visited.contains(root)) return
    visited.add(root)
    if (distance == k) {
        results.add(root.`val`)
        return
    }
    dfs(root.left, distance + 1, k)
    dfs(root.right, distance + 1, k)
    dfs(parentMap[root], distance + 1, k)
}

// Or equivalently, we track the parent node without using visited set.
private fun dfs(root: TreeNode?, source: TreeNode?, distance: Int, k: Int) {
    if (root == null) return 
    if (distance == k) {
        results.add(root.`val`)
        return
    }
    if (root.left != source) dfs(root.left, root, distance + 1, k)
    if (root.right != source) dfs(root.right, root, distance + 1, k)
    if (parentMap[root] != source) dfs(parentMap[root], root, distance + 1, k)
}

• Time Complexity: O(n) where n is the number of nodes in the tree.
• Space Complexity: O(n) for the parent map and visited set.