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

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

给定一个二叉树（具有根结点 root）， 一个目标结点 target ，和一个整数值 k 。

返回到目标结点 target 距离为 k 的所有结点的值的列表。 答案可以以 任何顺序 返回。

示例 1：

输入：root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, k = 2
输出：[7,4,1]
解释：所求结点为与目标结点（值为 5）距离为 2 的结点，值分别为 7，4，以及 1

method: BFS

将二叉树转化为图，然后用BFS

static const int N = 501;
static const int M = N * 2; 
struct edge {
    int u, v;
    int next;
}e[2 * M];
int head[N];
int cnt;
bool used[N];
void add(int u, int v) {
    e[++cnt].u = u;
    e[cnt].v = v;
    e[cnt].next = head[u];
    head[u] = cnt;
}
void dfs(TreeNode *root) {
    if (!root) return;
    if (root->left) {
        add(root->val, root->left->val);
        add(root->left->val, root->val);
        dfs(root->left);
    }
    if (root->right) {
        add(root->val, root->right->val);
        add(root->right->val, root->val);
        dfs(root->right);
    }
}
vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
    vector<int> res;
    dfs(root);
    queue<int> q;
    q.push(target->val);
    used[target->val] = true;
    int depth = 0;
    while (!q.empty()) {
        int size = q.size();
        for (int i = 0; i < size; i++) {
            int x = q.front();
            q.pop();
            if (depth == k) res.push_back(x);   // 获取深度为k的节点
            for (int i = head[x]; i; i = e[i].next) {
                int v = e[i].v;
                if (!used[v]) {
                    used[v] = true;
                    q.push(v); 
                }
            }
        }
        if (depth == k) break;  // 深度为k的节点都保存了，可以退出了
        depth++;
    }
    return res;
}