JAVA程序设计：二叉树中所有距离为 K 的结点（LeetCode：863）

给定一个二叉树（具有根结点 root）， 一个目标结点 target ，和一个整数值 K 。

返回到目标结点 target 距离为 K 的所有结点的值的列表。 答案可以以任何顺序返回。

 

示例 1：

输入：root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, K = 2

输出：[7,4,1]

解释：
所求结点为与目标结点（值为 5）距离为 2 的结点，
值分别为 7，4，以及 1

注意，输入的 "root" 和 "target" 实际上是树上的结点。
上面的输入仅仅是对这些对象进行了序列化描述。
 

提示：

给定的树是非空的，且最多有 K 个结点。
树上的每个结点都具有唯一的值 0 <= node.val <= 500 。
目标结点 target 是树上的结点。
0 <= K <= 1000.

思路：我的方法是首先遍历一波树，存储每个节点对应的父节点，之后我们可以从target节点开始层次遍历，他的子节点包括原始的左右节点以及之前预处理出来的父节点，然后找到距离该点为K的所有节点即可。

class Solution {
	
	TreeNode start=null;
	List<Integer> ans=new ArrayList<>();
	Map<Integer,TreeNode> map=new HashMap<>(); 
	
    public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
        
    	if(root==null || target==null) 
    		return ans;
    	
    	if(K==0) {
    		ans.add(target.val);
    		return ans;
    	}
    	
    	dfs(root,null,target.val);
    	
    	find(start.left,start.val,K-1);
    	find(start.right,start.val,K-1);
    	find(map.get(target.val),start.val,K-1);
    	
    	return ans;
    }
    
    private void find(TreeNode root,int f,int depth) {
    	
    	if(root==null) return;
    	if(depth==0) {
    		ans.add(root.val);
    		return;
    	}
    	if(root.left!=null && root.left.val!=f) 
    		find(root.left,root.val,depth-1);
    	if(root.right!=null && root.right.val!=f) 
    		find(root.right,root.val,depth-1);
    	if(map.get(root.val)!=null && map.get(root.val).val!=f)
    		find(map.get(root.val),root.val,depth-1);
    	
    }
    
    private void dfs(TreeNode root,TreeNode father,int target) {
    	
    	if(root==null) return;
    	if(root.val==target) start=root;
    	
    	map.put(root.val, father);
    	
    	dfs(root.left,root,target);
    	dfs(root.right,root,target);
    	
    }
}