编程训练第八十四期——二叉树中所有距离为 K 的结点

编程训练：

给定一个二叉树（具有根结点 root）， 一个目标结点 target ，和一个整数值 K 。返回到目标结点 target 距离为 K 的所有结点的值的列表。 答案可以以任何顺序返回。

---

示例：

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

---

解法：

1.DFS+BFS
首先使用 DFS 遍历二叉树，为每个节点保存其父结点。再使用 BFS 向三个方向搜索（左孩子、右孩子、父结点），当搜索第 K 次时，队列中的全部结点即为所求。
时间复杂度O(N)
空间复杂度O(N)

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
private:
    void dfs(TreeNode *root, TreeNode *pre)
    {
        if (root)
        {
            parent.insert({root, pre});
            dfs(root->left, root);
            dfs(root->right, root);
        }
    }
public:
    unordered_map<TreeNode *, TreeNode *> parent;
    vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {
        dfs(root, NULL);
        queue<TreeNode *> que;
        que.push(target);
        unordered_set<TreeNode *> visited;
        visited.insert(target);
        while (!que.empty())
        {
            if (K-- == 0)
            {
                vector<int> res;
                while (!que.empty())
                {
                    res.push_back(que.front()->val);
                    que.pop();
                }
                return res;
            }
            else
            {
                int size = que.size();
                for (int i = 0; i < size; i++)
                {
                    TreeNode *node = que.front();
                    que.pop();
                    if (node->left && visited.find(node->left) == visited.end())
                    {
                        que.push(node->left);
                        visited.insert(node->left);
                    }

                    if (node->right && visited.find(node->right) == visited.end())
                    {
                        que.push(node->right);
                        visited.insert(node->right);
                    }

                    if (parent[node] && visited.find(parent[node]) == visited.end())
                    {
                        que.push(parent[node]);
                        visited.insert(parent[node]);
                    }
                }
            }
        }
        return vector<int> {};
    }
};