prep-coding/leetcode/DistanceK.java at master · tushar-sharma/prep-coding · 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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {

    public static void buildParentMap(TreeNode curr, Map<TreeNode, TreeNode> parent) {

           if (curr == null){
               return;
           }

           if (curr.left != null) {
               parent.put(curr.left, curr);
           }

           if (curr.right != null) {
               parent.put(curr.right, curr);
           }

           buildParentMap(curr.left, parent);

           buildParentMap(curr.right, parent);

    }

    public static void findNodes(TreeNode target,
                                 int k,
                                 int accumulator,
                                 Map<TreeNode, TreeNode> parent,
                                 List<Integer> result,
                                 Set<TreeNode> visited){

        if (target == null || accumulator > k) {
            return;
        }

        if (k == accumulator){
            result.add(target.val);
        }

        visited.add(target);

        if (!visited.contains(target.left)) {
            findNodes(target.left,
                  k,
                 accumulator + 1,
                 parent,
                 result,
                 visited);

        }


        if (!visited.contains(target.right)) {

            findNodes(target.right,
                  k,
                  accumulator + 1,
                 parent,
                 result,
                 visited);
        }

        if (!visited.contains(parent.get(target))) {

            findNodes(parent.get(target),
                   k,
                   accumulator + 1,
                   parent,
                  result,
                  visited);
        }


    }
    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {


        Map<TreeNode, TreeNode> parent = new HashMap<>();

        List<Integer> result = new ArrayList<>();

        buildParentMap(root, parent);

        Set<TreeNode> visited = new HashSet<>();

        findNodes(target,
                  k,
                  0,
                  parent,
                  result,
                  visited);

        return result;

    }
}