863. All Nodes Distance K in Binary Tree

Problem Link

Description

---

Given the root of a binary tree, the value of a target node target, and an integer k, return an array of the values of all nodes that have a distance k from the target node.

You can return the answer in any order.

 

Example 1:

Input: root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, k = 2
Output: [7,4,1]
Explanation: The nodes that are a distance 2 from the target node (with value 5) have values 7, 4, and 1.

Example 2:

Input: root = [1], target = 1, k = 3
Output: []

 

Constraints:

• The number of nodes in the tree is in the range [1, 500].
• 0 <= Node.val <= 500
• All the values Node.val are unique.
• target is the value of one of the nodes in the tree.
• 0 <= k <= 1000

Solution

---

Python3

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None
 
class Solution:
    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:
        res = []
        mp = {}
 
        def dfs1(node):
            if not node: return -1
 
            if node == target:
                mp[node] = 0
                return 0
            
            if (left := dfs1(node.left)) != -1:
                mp[node] = left + 1
                return left + 1
            
            if (right := dfs1(node.right)) != -1:
                mp[node] = right + 1
                return right + 1
            
            return -1
        
        def dfs2(node, length):
            if not node: return
 
            if node in mp:
                length = mp[node]
            
            if length == k:
                res.append(node.val)
            
            dfs2(node.left, length + 1)
            dfs2(node.right, length + 1)
 
        dfs1(root)
        dfs2(root, mp[root])
 
        return res