[LeetCode] 863. All Nodes Distance K in Binary Tree

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

It is straightforward  to use BFS to collect the nodes on distance K in the binary tree. 

To support move upward and downward simultaneously, we need to add one more pointer to point from each node to its parent node. So we convert the given binary tree to a directed graph. 

# 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]:
        
        # use dfs to find parent of each nodes
        parent = {root: None}
        def dfs(node):
            for child in [node.left, node.right]:
                if child:
                    parent[child] = node
                    dfs(child)
        
        dfs(root)
        
        # start from target node and use bfs to find all nodes on distance k
        queue = deque([target])
        visited = set([target])
        
        result = []
        step = 0
        
        while queue and step <= k:
            for _ in range(len(queue)):
                node = queue.popleft()
                if step == k:
                    result.append(node.val)
                
                for nxt in [node.left, node.right, parent[node]]:
                    if nxt and nxt not in visited:
                        visited.add(nxt)
                        queue.append(nxt)
                
            step += 1
                
        return result