// 863. All Nodes Distance K in Binary Tree #include #include /* * given the 'root' of a binary tree, the value of a target node 'target', and * an integer 'k', return an array of the values of all nodes that have a * distance 'k' from the target node. you can return the answer in any order. */ struct TreeNode { int val; struct TreeNode *left; struct TreeNode *right; }; int depth(struct TreeNode *root) { if (!root) return -1; return fmax(depth(root->left), depth(root->right)) + 1; } void find_path(struct TreeNode *root, struct TreeNode *target, struct TreeNode **data, int data_idx, struct TreeNode **path, int *path_idx) { if (!root) return; data[data_idx] = root; if (root == target) { *path_idx = data_idx + 1; for (int i = 0; i < data_idx + 1; i++) path[i] = data[i]; return; } find_path(root->left, target, data, data_idx + 1, path, path_idx); find_path(root->right, target, data, data_idx + 1, path, path_idx); } void k_level_data(struct TreeNode *root, int *data, int *idx, int k) { if (!root) return; if (!k) { data[*idx] = root->val; (*idx)++; return; } k_level_data(root->left, data, idx, k - 1); k_level_data(root->right, data, idx, k - 1); } int *distanceK(struct TreeNode *root, struct TreeNode *target, int k, int *return_size) { int *ans = malloc(100 * sizeof(int)), idx = 0; if (root == target) { k_level_data(root, ans, &idx, k); *return_size = idx; return ans; } int dep = depth(root); struct TreeNode **data = malloc((dep + 1) * sizeof(struct TreeNode *)); struct TreeNode **path = malloc((dep + 1) * sizeof(struct TreeNode *)); int data_idx = 0, path_idx = 0, a = 1; find_path(root, target, data, data_idx, path, &path_idx); k_level_data(target, ans, &idx, k); for (int i = path_idx - 2; i >= 0; i--) { if (!(k - a)) { ans[idx] = path[i]->val; idx++; break; } if (path[i]->left == path[i + 1]) k_level_data(path[i]->right, ans, &idx, k - a - 1); else k_level_data(path[i]->left, ans, &idx, k - a - 1); a++; } *return_size = idx; return ans; }