1#include <stdio.h>
2#include <stdlib.h>
3#include <stdbool.h>
4#include <string.h>
5
6// Define the TreeNode structure
7struct TreeNode {
8    int val;
9    struct TreeNode *left;
10    struct TreeNode *right;
11};
12
13// A node in the queue
14struct QueueNode {
15    struct TreeNode *node;
16    int distance;
17};
18
19// Create a new tree node
20struct TreeNode* newTreeNode(int val) {
21    struct TreeNode* node = (struct TreeNode*)malloc(sizeof(struct TreeNode));
22    node->val = val;
23    node->left = node->right = NULL;
24    return node;
25}
26
27// A simple queue implementation
28#define MAX_QUEUE_SIZE 1000
29typedef struct {
30    struct QueueNode data[MAX_QUEUE_SIZE];
31    int front, rear;
32} Queue;
33
34void initQueue(Queue *q) {
35    q->front = q->rear = 0;
36}
37
38bool isQueueEmpty(Queue *q) {
39    return q->front == q->rear;
40}
41
42void enqueue(Queue *q, struct TreeNode *node, int distance) {
43    q->data[q->rear].node = node;
44    q->data[q->rear].distance = distance;
45    q->rear++;
46}
47
48struct QueueNode dequeue(Queue *q) {
49    return q->data[q->front++];
50}
51
52// Helper function for DFS that creates a map from node value to parent node
53void populateParentMap(struct TreeNode *root, struct TreeNode *parent, struct TreeNode **parentMap) {
54    if (root == NULL) return;
55    parentMap[root->val] = parent;
56    populateParentMap(root->left, root, parentMap);
57    populateParentMap(root->right, root, parentMap);
58}
59
60// BFS to find all nodes at distance K
61void findNodesAtDistance(struct TreeNode *target, int K, struct TreeNode **parentMap, int *result, int *returnSize) {
62    Queue q;
63    initQueue(&q);
64    bool visited[501] = {false};
65    enqueue(&q, target, 0);
66    visited[target->val] = true;
67
68    while (!isQueueEmpty(&q)) {
69        struct QueueNode current = dequeue(&q);
70        struct TreeNode *currNode = current.node;
71        int currDistance = current.distance;
72
73        if (currDistance == K) {
74            result[(*returnSize)++] = currNode->val;
75        }
76
77        if (currDistance > K) break;
78
79        // Check left child
80        if (currNode->left && !visited[currNode->left->val]) {
81            enqueue(&q, currNode->left, currDistance + 1);
82            visited[currNode->left->val] = true;
83        }
84
85        // Check right child
86        if (currNode->right && !visited[currNode->right->val]) {
87            enqueue(&q, currNode->right, currDistance + 1);
88            visited[currNode->right->val] = true;
89        }
90
91        // Check parent
92        struct TreeNode *parent = parentMap[currNode->val];
93        if (parent && !visited[parent->val]) {
94            enqueue(&q, parent, currDistance + 1);
95            visited[parent->val] = true;
96        }
97    }
98}
99
100int* distanceK(struct TreeNode* root, struct TreeNode* target, int K, int* returnSize) {
101    struct TreeNode *parentMap[501] = {NULL};
102    int *result = malloc(500 * sizeof(int));
103    *returnSize = 0;
104
105    populateParentMap(root, NULL, parentMap);
106    findNodesAtDistance(target, K, parentMap, result, returnSize);
107
108    return result;
109}
110

1#include <iostream>
2#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;
}