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
| class Solution {
public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
Map<TreeNode, List<TreeNode>> nodeToNeighborList = new HashMap<>();
dfs(root, null, nodeToNeighborList);
List<Integer> numList = new ArrayList<>();
Set<TreeNode> visitedNodeSet = new HashSet<>();
dfs(target, k, nodeToNeighborList, visitedNodeSet, numList);
return numList;
}
private void dfs(TreeNode target, int k, Map<TreeNode, List<TreeNode>> nodeToNeighborList, Set<TreeNode> visitedNodeSet, List<Integer> numList) {
if (visitedNodeSet.contains(target)) {
return;
}
visitedNodeSet.add(target);
if (k == 0) {
numList.add(target.val);
return;
}
for (TreeNode neighbor : nodeToNeighborList.getOrDefault(target, Collections.emptyList())) {
dfs(neighbor, k - 1, nodeToNeighborList, visitedNodeSet, numList);
}
}
private void dfs(TreeNode node, TreeNode parent, Map<TreeNode, List<TreeNode>> nodeToNeighborList) {
if (node == null) {
return;
}
List<TreeNode> neighbors = nodeToNeighborList.computeIfAbsent(node, k -> new ArrayList<>());
if (parent != null) {
neighbors.add(parent);
}
if (node.left != null) {
neighbors.add(node.left);
dfs(node.left, node, nodeToNeighborList);
}
if (node.right != null) {
neighbors.add(node.right);
dfs(node.right, node, nodeToNeighborList);
}
}
}
|