[leetcode]863. All Nodes Distance K in Binary Tree

Analysis

normal day—— [每天刷題並不難0.0]

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.

Implement

/**
 * 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 {
public:
    vector<int> distanceK(
 
TreeNode* root, TreeNode* target, int K) {
        if(!root)
            return res;
        findParent(root);
        dfs(target, K);
        return res;
    }
    void findParent(TreeNode* node){
        if(!node)
            return;
        if(node->left){
            parent[node->left] = node;
 

            findParent(node->left);
        }
        if(node->right){
            parent[node->right] = node;
            findParent(node->right);
        }
    }
    void dfs(TreeNode* node, int k){
        if(visit.find(node) != visit.end())
            return ;
        visit.insert(node);
        if(k == 0)
            res.push_back(node->val);
        if(node->left)
            dfs(node->left, k-1);
        if(node->right)
            dfs(node->right, k-1);
        TreeNode* p = parent[node];
        if(p)
            dfs(p, k-1);
    }
private:
    vector<int> res;
    map<TreeNode*, TreeNode*> parent;
    set<TreeNode*> visit;
};