算法：与 target的距离是K的所有节点

题目描述

给定三个参数:

• 二叉树的头节点head,树上某个节点 target,正数K
• 从 target开始，可以向上走或者向下走
• 返回与 target的距离是K的所有节点

struct TreeNode {
    int val;
    TreeNode *left;
    TreeNode *right;
    TreeNode() : val(0), left(nullptr), right(nullptr) {}
    TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
    TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
};
class Solution {
    
public:
    std::vector<TreeNode *> distance(TreeNode *root, TreeNode *target, int k){
        
    }
};

题目解析

思路：

• 生成每个节点的父节点map，这棵树中每一个节点都能找到它的父节点是谁： 因为经典二叉树中没有往上左的职责，父map的功能是让一个节点往上左
• 当有这个父map之后，距离一个节点为k的路径就变成了图的宽度优先遍历

实现：

class Solution {
    void createParentMap(TreeNode * curr, std::map<TreeNode *, TreeNode *> parents){
        if(curr == nullptr){
            return;
        }
        
        if(curr->left != nullptr){
            parents[curr->left] = curr;
            createParentMap(curr->left, parents);
        }
        
        if(curr->right != nullptr){
            parents[curr->right] = curr;
            createParentMap(curr->right, parents);
        }
    }
public:
    std::vector<TreeNode *> distance(TreeNode *root, TreeNode *target, int k){
        if(root == nullptr){
            return {};
        }
        std::map<TreeNode *, TreeNode *> parents;
        parents[root] = nullptr;
        createParentMap(root, parents);
        
        std::queue<TreeNode*> queue;
        std::set<TreeNode *> visited;
        queue.emplace(target);
        visited.emplace(target);
        int currLevel = 0;
        std::vector<TreeNode *> ans;
        while (!queue.empty()){
            int size = queue.size();
            while (size-- > 0){
                auto cur = queue.front(); queue.pop();
                if(currLevel == k){
                    ans.emplace_back(cur);
                }
                
                // 它的左右孩子&父节点均是它的下一层
                if(cur->left != nullptr && !visited.count(cur->left)){
                    visited.emplace(cur->left);
                    queue.emplace(cur->left);
                }
                if(cur->right != nullptr && !visited.count(cur->right)){
                    visited.emplace(cur->right);
                    queue.emplace(cur->right);
                }
                if(parents.count(cur) && !visited.count(parents[cur])){
                    visited.emplace(parents[cur]);
                    queue.emplace(parents[cur]);
                }
            }

			// 当前层搞完了
            currLevel++;
            if(currLevel > k){
                break;
            }
        }
        
        return ans;
    }
};