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

1#include <iostream>
#include <vector>
#include <unordered_map>
#include <unordered_set>
using namespace std;

struct TreeNode {
    int val;
    TreeNode* left;
    TreeNode* right;
    TreeNode(int x) : val(x), left(NULL), right(NULL) {}
};

void findChildrenAtDistance(TreeNode* node, int K, vector<int>& result, unordered_set<TreeNode*>& visited) {
    if (!node || visited.count(node)) return;
    if (K == 0) {
        result.push_back(node->val);
        return;
    }
    visited.insert(node);
    findChildrenAtDistance(node->left, K - 1, result, visited);
    findChildrenAtDistance(node->right, K - 1, result, visited);
    visited.erase(node);
}

bool findTargetAndParents(TreeNode* root, TreeNode* target, unordered_map<TreeNode*, TreeNode*>& parentMap) {
    if (!root) return false;
    if (root == target) return true;
    parentMap[root->left] = root;
    if (findTargetAndParents(root->left, target, parentMap)) return true;
    parentMap[root->right] = root;
    if (findTargetAndParents(root->right, target, parentMap)) return true;
    return false;
}

vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {
    unordered_map<TreeNode*, TreeNode*> parentMap;
    vector<int> result;
    unordered_set<TreeNode*> visited;
    findTargetAndParents(root, target, parentMap);

    queue<pair<TreeNode*, int>> bfsQueue;
    bfsQueue.push({target, 0});
    visited.insert(target);

    while (!bfsQueue.empty()) {
        auto [node, dist] = bfsQueue.front(); bfsQueue.pop();

        if (dist == K) result.push_back(node->val);

        for (TreeNode* neighbor : {node->left, node->right, parentMap[node]}) {
            if (neighbor && !visited.count(neighbor)) {
                visited.insert(neighbor);
                bfsQueue.push({neighbor, dist + 1});
            }
        }
    }

    return result;
}