LeetCode - 863. All Nodes Distance K in Binary Tree

We are given a binary tree (with root node root), a target node, and an integer value K.

Return a list of the values of all nodes that have a distance K from the target node.  The answer can be returned in any order.

Example 1:

Input: root = [3, 5, 1, 6, 2, 0, 8, null, null, 7, 4], target = 5, K = 2

Output: [7, 4, 1]

Explanation:

The nodes that are a distance 2 from the target node (with value 5) have values 7, 4, and 1.

思路：

    由于不仅包括target的子树，还包括其他部分，所以必须要知道parent信息，如果树结构中存储了parent信息就不需要在遍历然后set_parent了（set_parent 对 LeetCode 中很多树的题是很有用的，不知道为什么树结构中没有父节点信息）。有了parent信息之后，在target的子树中以及target的parent中递归，每递归一层，K就少1，直到为0。dfs很容易解决。

class Solution {
public:
    vector<int> res;
    map<TreeNode*, TreeNode*> parent;
    set<TreeNode*> visited;
    
    void dfs(TreeNode* tn, int K)
    {
        if(visited.find(tn) != visited.end())
            return;
        
        visited.insert(tn);
        if(K == 0)
            res.push_back(tn->val);
        else
        {
            if(tn->left)
                dfs(tn->left, K - 1);
            if(tn->right)
                dfs(tn->right, K - 1);
        }
        TreeNode* p = parent[tn];
        if(p)
            dfs(p, K - 1);
        return;
    }
    
    void set_parent(TreeNode* tn)
    {
        if(tn->left)
        {
            parent[tn->left] = tn;
            set_parent(tn->left);
        }
        if(tn->right)
        {
            parent[tn->right] = tn;
            set_parent(tn->right);
        }
        return;
    }
    
    vector<int> distanceK(TreeNode* root, TreeNode* target, int K)
    {
        if(root == NULL)
            return vector<int>();
        set_parent(root);   // 记录所有节点的父信息
        dfs(target, K);
        return res;
    }
}；