1import java.util.*;
2
3class TreeNode {
4    int val;
5    TreeNode left;
6    TreeNode right;
7    TreeNode(int x) { val = x; }
8}
9
10public class Solution {
11    public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
12        Map<TreeNode, TreeNode> parentMap = new HashMap<>();
13        buildParentMap(root, parentMap);
14
15        Queue<TreeNode> queue = new LinkedList<>();
16        Set<TreeNode> visited = new HashSet<>();
17        queue.add(target);
18        visited.add(target);
19
20        int currentLevel = 0;
21        while (!queue.isEmpty()) {
22            if (currentLevel == K) {
23                List<Integer> result = new ArrayList<>();
24                for (TreeNode node : queue) {
25                    result.add(node.val);
26                }
27                return result;
28            }
29            int size = queue.size();
30            for (int i = 0; i < size; i++) {
31                TreeNode node = queue.poll();
32                if (node.left != null && !visited.contains(node.left)) {
33                    queue.add(node.left);
34                    visited.add(node.left);
35                }
36                if (node.right != null && !visited.contains(node.right)) {
37                    queue.add(node.right);
38                    visited.add(node.right);
39                }
40                if (parentMap.containsKey(node) && !visited.contains(parentMap.get(node))) {
41                    queue.add(parentMap.get(node));
42                    visited.add(parentMap.get(node));
43                }
44            }
45            currentLevel++;
46        }
47
48        return new ArrayList<>();
49    }
50
51    private void buildParentMap(TreeNode node, Map<TreeNode, TreeNode> parentMap) {
52        if (node.left != null) {
53            parentMap.put(node.left, node);
54            buildParentMap(node.left, parentMap);
55        }
56        if (node.right != null) {
57            parentMap.put(node.right, node);
58            buildParentMap(node.right, parentMap);
59        }
60    }
61}
62

1using System;
2using System.Collections.Generic;

public class TreeNode {
    public int val;
    public TreeNode left;
    public TreeNode right;
    public TreeNode(int x) { val = x; }
}

public class Solution {
    public IList<int> DistanceK(TreeNode root, TreeNode target, int K) {
        var parentMap = new Dictionary<TreeNode,TreeNode>();
        BuildParentMap(root, null, parentMap);

        var result = new List<int>();
        var visited = new HashSet<TreeNode>();
        Dfs(target, K, null, visited, parentMap, result);
        return result;
    }

    private void BuildParentMap(TreeNode node, TreeNode parent, Dictionary<TreeNode,TreeNode> parentMap) {
        if (node != null) {
            parentMap[node] = parent;
            BuildParentMap(node.left, node, parentMap);
            BuildParentMap(node.right, node, parentMap);
        }
    }

    private void Dfs(TreeNode node, int K, TreeNode parent, HashSet<TreeNode> visited, Dictionary<TreeNode,TreeNode> parentMap, IList<int> result) {
        if (node == null || visited.Contains(node)) return;
        visited.Add(node);
        if (K == 0) {
            result.Add(node.val);
            return;
        }
        if (node.left != parent) Dfs(node.left, K - 1, node, visited, parentMap, result);
        if (node.right != parent) Dfs(node.right, K - 1, node, visited, parentMap, result);
        TreeNode parentNode = parentMap[node];
        if (parentNode != null && parentNode != parent) Dfs(parentNode, K - 1, node, visited, parentMap, result);
    }
}