给定一个二叉树(具有根结点 root), 一个目标结点 target ,和一个整数值 K 。
返回到目标结点 target 距离为 K 的所有结点的值的列表。 答案可以以任何顺序返回。
代码语言:javascript复制class Solution {
// 用map记录每个节点的父节点
private Map<TreeNode, TreeNode> parents = new HashMap<>();
private Set<TreeNode> used = new HashSet<>();
private TreeNode targetNode;
List<Integer> res = new LinkedList<>();
public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
find(root, null, target);
dfs(targetNode, K);
return res;
}
private void find(TreeNode root, TreeNode parent, TreeNode target) {
if (null == root) {
return;
}
if (root.val == target.val) {
targetNode = root;
}
parents.put(root, parent);
find(root.left, root, target);
find(root.right, root, target);
}
// 找到目标节点后以目标节点为开始位置向三个方向蔓延
private void dfs(TreeNode root, int distance) {
if (root != null && !used.contains(root)) {
// 标记为已访问
used.add(root);
if (distance <= 0) {
res.add(root.val);
return;
}
dfs(root.left, distance - 1);
dfs(root.right, distance - 1);
dfs(parents.get(root), distance - 1);
}
}
}