863. All Nodes Distance K in Binary Tree

이진 트리의 노드에서 거리가 k인 모든 노드의 값 배열을 반환

class Solution {
        
    Map<TreeNode, Integer> map = new HashMap<>();
    List<Integer> res = new LinkedList<>();
        
    // target 노드에서 k 거리에 있는 모든 노드의 값 배열
    public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
        find(root, target);
        dfs(root, target, K, map.get(root));
        return res;
    }
    
    // find target node first and store the distance in that path that we could use it later directly
    private int find(TreeNode root, TreeNode target) {
        if (root == null) return -1;
        
        if (root == target) {
            map.put(root, 0);
            return 0;
        }
        
        int left = find(root.left, target);
        if (left >= 0) {
            map.put(root, left + 1);
            return left + 1;
        }
        
		int right = find(root.right, target);
		if (right >= 0) {
            map.put(root, right + 1);
            return right + 1;
        }
        
        return -1;
    }
    
    private void dfs(TreeNode root, TreeNode target, int K, int length) {
        if (root == null) return;
        
        if (map.containsKey(root)) {
            length = map.get(root);
        }
        
        if (length == K) {
            res.add(root.val);
        }
        
        dfs(root.left, target, K, length + 1);
        dfs(root.right, target, K, length + 1);
    }
}






© 2017. by yeopoong.github.io

Powered by yeopoong