leetcode——863. 二叉树中所有距离为 K 的结点

最新推荐文章于 2024-08-23 14:25:42 发布

原创最新推荐文章于 2024-08-23 14:25:42 发布 · 213 阅读

0 ·

CC 4.0 BY-SA版权

Leetcode 专栏收录该内容

145 篇文章

订阅专栏

本文介绍了如何使用深度优先和广度优先搜索算法在二叉树中查找深度为K的所有节点，通过递归构建父节点索引，并利用队列避免死循环。关键步骤包括构建节点父子关系、初始化队列和访问记录，以及在广度优先搜索过程中更新深度和输出结果。

思路——深度遍历+广度遍历

建立结点的父亲索引——通过递归实现
从target出发，进行广度搜索查找depth为K的所有结点
需要建立访问记录，防止死循环（因为遍历会向上也会向下走）

代码

步骤：
1. 建立父亲索引
2. 将target放入队列
3. 对target左右孩子、父亲结点进行广度搜索
4. 找到depth 进行队列全弹出
5. 输出答案

class Solution {
private:
    // 递归方式来构建父节点信息
    unordered_map<TreeNode*, TreeNode*> node2parent;
    void dfs(TreeNode* curr, TreeNode* parent){
        if (curr != nullptr){
            node2parent[curr] = parent;
            dfs(curr->left, curr);
            dfs(curr->right, curr);
        }       
    }

public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {
        // 构建好映射表
        dfs(root, nullptr);

        // target放入队列
        queue<TreeNode*> q;
        unordered_set<TreeNode*> visited;
        q.push(target);
        visited.insert(target);
        
        //广度搜索
        int dist = 0;
        while (!q.empty()){
            // 检查是否满足条件
            if (dist == K){
                vector<int> res;
                while (!q.empty()){
                    res.emplace_back(q.front()->val);
                    q.pop();
                }
                return res;
            }

            // 保证每次只遍历当前的层的数量
            for (int i = q.size(); i > 0; --i){
                TreeNode* curr = q.front();q.pop();
                // 从三个方向去遍历
                if (curr->left != nullptr && visited.find(curr->left) == visited.end()){
                    q.push(curr->left);
                    visited.insert(curr->left);
                }
                if (curr->right != nullptr && visited.find(curr->right) == visited.end()){
                    q.push(curr->right);
                    visited.insert(curr->right);
                }
                TreeNode* parent = node2parent[curr];
                if (parent!=nullptr && visited.find(parent) == visited.end()) {
                    q.push(parent);
                    visited.insert(parent);
                }
            }
            // 一层遍历完成距离+1
            ++dist;
        }
        
        // 默认返回空
        return {};
    }
};