代码改变世界

863 二叉树中所有距离为 K 的结点

2021-12-01 00:14  tonyniu8  阅读(37)  评论(0编辑  收藏  举报

思路:
1.先dfs記錄parent,
2.然後bfs找到第k層,其中要注意,已經便利過的不能重複變臉,所以要一個seen的集合來保存已經便利過的。
3.注意java stream的用法。
https://www.runoob.com/java/java8-streams.html

queue.stream().map(n -> n.val).collect(Collectors.toList());
class Solution {
     	    private Map<TreeNode, TreeNode>parents = new HashMap<>();
	    Queue<TreeNode>queue = new LinkedList<>();
	    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
	       dfs(root);

	       queue.offer(target);
	       return bfs(target, k);
	    }
	    private void dfs(TreeNode node){
	        if(node.left != null){
	            parents.put(node.left,node);
	            dfs(node.left);
	        }
	        if(node.right != null){
	            parents.put(node.right,node);
	            dfs(node.right);
	        }
	    }
	    private List<Integer> bfs(TreeNode target, int k){
	        int depth=0;
	        Set<TreeNode>seen = new HashSet<>();
	        seen.add(target);
	        while(!queue.isEmpty()&&depth<k){
	            int levelSize=queue.size();
	            depth++;
	            for (int i=0;i<levelSize;i++){
	                TreeNode node = queue.poll();
	                if(node.left != null&&!seen.contains(node.left)){
	                    queue.offer(node.left);
	                    seen.add(node.left);
	                }
	                if(node.right != null&&!seen.contains(node.right)){
	                    queue.offer(node.right);
	                    seen.add(node.right);
	                }
	                TreeNode parent=parents.get(node);
	                if(parent!= null&&!seen.contains(parent)){
	                    queue.offer(parent);
	                    seen.add(parent);
	                }
	            }
	        }
	        return queue.stream().map(node->node.val).collect(Collectors.toList());
	    }
}