Java/863.all-nodes-distance-k-in-binary-tree.java at master · vivek880490/Java · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
/*
 * @lc app=leetcode id=863 lang=java
 *
 * [863] All Nodes Distance K in Binary Tree
 */

// @lc code=start
/**
 * 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) {
        ArrayList<Integer> al = new ArrayList<>();
        if (root == null)
            return al;
        HashMap<TreeNode, TreeNode> map = new HashMap<>();
        parent(map, root);
        HashSet<TreeNode> visited = new HashSet<>();
        Queue<TreeNode> q = new LinkedList<>();
        q.add(target);

        while (!q.isEmpty()) {
            int size = q.size();
            for (int i = 0; i < size; i++) {
                TreeNode cur = q.remove();
                visited.add(cur);
                if (k == 0) {
                    al.add(cur.val);
                }

                if (map.containsKey(cur) && !visited.contains(map.get(cur))) {
                    q.add(map.get(cur));
                }
                if (cur.left != null && !visited.contains(cur.left)) {
                    q.add(cur.left);
                }

                if (cur.right != null && !visited.contains(cur.right)) {
                    q.add(cur.right);
                }
            }
            k--;
            if (k < 0)
                break;
        }
        return al;

    }

    void parent(HashMap<TreeNode, TreeNode> map, TreeNode root) {

        if (root == null)
            return;
        if (root.left != null) {
            map.put(root.left, root);
        }

        if (root.right != null) {
            map.put(root.right, root);
        }

        parent(map, root.left);
        parent(map, root.right);
        return;
    }

}
// @lc code=end