All nodes distance K binary tree

gistfile1.txt

public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
    Map<Integer, Integer> map = new HashMap<>();  // distance to target from root
    find(root, target, map);

    System.out.println("Map");
    for (Map.Entry<Integer, Integer> entry : map.entrySet()) {
        System.out.println(entry.getKey() + " ->" + entry.getValue());
    }
    List<Integer> res = new ArrayList<>();
    helper(root, target, K, map.get(root.val), map, res);
    return res;
}
private int find(TreeNode root, TreeNode target, Map<Integer, Integer> map) {

    if (root == null) return -1;
    System.out.println("root: " + root.val);

    if (root == target) {
        map.put(root.val, 0);
        return 0;
    }
    System.out.println("pass root == target");
    int left = find(root.left, target, map);
    if (left >= 0) {
        map.put(root.val, left + 1);  // miss, 这里要 +1
        return left + 1;
    }
    int right = find(root.right, target, map);
    if (right >= 0) {
        map.put(root.val, right + 1);
        return right + 1;
    }
    return -1;
}
private void helper(TreeNode root, TreeNode target, int K, int length, 
                    Map<Integer, Integer> map, List<Integer> res) {
    if (root == null) return;
    // use the length in map if there's such node in map
    if (map.containsKey(root.val)) {
        length = map.get(root.val);
    }
    if (length == K) res.add(root.val);
    helper(root.left, target, K, length + 1, map, res);
    helper(root.right, target, K, length + 1, map, res);
}