863. 二叉树中所有距离为 K 的结点

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);
        }
    }
}