Leetcode 863. All nodes with distance K in the binary tree (the more difficult binary tree problem)

 

The most direct way of thinking is to convert the binary tree into a graph, and then use BFS to find the node with a distance of K. This is actually a new map based on the original tree.

/**
 * 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) {
        buildGraph(root);
        queue<TreeNode*> q;
        vector<int> res;
        unordered_set<TreeNode*> visited;
        q.push(target);
        visited.insert(target);
        int count = 0;
        while(!q.empty()&&count<=K){
            int n = q.size();
            while(n--){
                auto tmp = q.front();
                q.pop();
                if(count==K) res.emplace_back(tmp->val);
                for(auto neighbor:graph[tmp]){
                    if(!visited.count(neighbor)){
                        q.push(neighbor);
                        visited.insert(neighbor);
                    }
                }
            }
            count++;
        }
        return res;
    }

    void buildGraph(TreeNode* root){
        if(root==NULL){
            return;
        }
        if(root->left){
            graph[root].emplace_back(root->left);
            graph[root->left].emplace_back(root);
        }
        if(root->right){
            graph[root].emplace_back(root->right);
            graph[root->right].emplace_back(root);
        }
        buildGraph(root->left);
        buildGraph(root->right);
    }
unordered_map<TreeNode*,vector<TreeNode*>> graph;

};

 

Using the characteristics of the tree is more difficult to think of with a recursive function. The idea is as shown in the figure. Assuming that we know that the distance between the left subtree B of root A and the Target is l, then the distance between the right subtree C and T is l+2. Nodes with a distance of K to the target can be divided into two categories, one is in the subtree below the target as the root node, and the other part is above. The upper node can be found by finding the node with a distance of KL-2 from the current root node. .

/**
 * 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) {
        dfs(root,target,K);
        return res;
    }

    vector<int> res;

    // 这个函数表示从根节点到target的距离
    int dfs(TreeNode* root, TreeNode* target, int K){
        if(root==nullptr){
            return -1;
        }
        if(root==target){
            collect(target,K);    // 从当前节点出发，将距离为K的节点加入答案
            return 0;
        }
        int left = dfs(root->left,target,K);
        int right = dfs(root->right,target,K);
        if(left>=0){
            if(left==K-1) res.push_back(root->val);
            collect(root->right,K-left-2);
            return left+1;
        }
        if(right>=0){
            if(right==K-1) res.push_back(root->val);
            collect(root->left,K-right-2);
            return right+1;
        }
        return -1;
    }

    void collect(TreeNode* root, int d){
        if(root==nullptr||d<0) return;
        if(d==0) res.push_back(root->val);
        collect(root->left,d-1);
        collect(root->right,d-1);
    }
};