LeetCode 863 All nodes distance k in binary tree

1. Problem Description

We are given a binary tree (with root node root), a target node, and an integer value K.

Return a list of the values of all nodes that have a distance K from the target node. The answer can be returned in any order.

Example

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.

2. My Thought About This Problem:

At first glance, this problem is a Tree Problem. However, think a little bit more, tree is a special kind of Undirected Graph. Each Node in tree has **at most three edges, one for left child, one for right child, and other one for its parent. **A small trick here is that the parent of root is root itself. The parent information can be stored in a Map when traversing the tree once.

After building all the relationships, the rest of work is to find Distance K nodes from the target node. This can be traversed level by level. This type of search is typically breath-first search. To prevent traverse back to the nodes we have traversed in previous step. We need another additional data structure which is the set to keep track of those nodes which have been traversed before.

You could pratice on Leetcode online IDE here.

Code

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {


    public List<Integer> distanceK(TreeNode root, TreeNode target, int K){
        Map<TreeNode, TreeNode> parentsMap = new HashMap<>();
        buildMap(root, root,parentsMap);
        return findKdistanceNode(target, K, parentsMap);    
    }

    /*
     * Traverse the tree in DFS way and record every node's parent. 
     */
    private static void buildMap(TreeNode node, TreeNode parent, Map<TreeNode, TreeNode> parentsMap){
        if(node == null) return;
        parentsMap.put(node,parent);
        if(node.left != null)  buildMap(node.left, node, parentsMap);
        if(node.right != null) buildMap(node.right, node, parentsMap);
    }

    /*
     * Given the target node and distance K, find the K distance Nodes. 
     */
    private static List<Integer> findKdistanceNode(TreeNode target, int K, Map<TreeNode, TreeNode> parentsMap){
        Set<TreeNode> traversed = new HashSet<>();
        Deque<TreeNode> queue  = new ArrayDeque<>();

        queue.offer(target);
        traversed.add(target);
        for(int level = 0; level< K; level++){

            int queueSize = queue.size();
            for(int idx =0 ; idx < queueSize; idx++){
                // traverse all its parents. 
                TreeNode curNode = queue.poll();
                traverseAndMark(curNode.left, traversed, queue);
                traverseAndMark(curNode.right, traversed, queue);
                TreeNode curNodeParent = parentsMap.get(curNode);
                traverseAndMark(curNodeParent, traversed, queue);
            }
        }

        //convert to the required result format.
        List<Integer> resultList = new ArrayList<>();

        queue.forEach(node -> {
            resultList.add(node.val);
        });

        return resultList;        
    }

    /*
     * Traverse those nodes which have not been visited and add to queue.
     */
    private static void traverseAndMark(TreeNode node, Set<TreeNode> traversed,Deque<TreeNode> queue){
        if(node != null && !traversed.contains(node)){
            traversed.add(node);
            queue.offer(node);
        }
    }
}