【刷题·二叉树·图】2385. 感染二叉树需要的总时间

前言

我觉得这个题目挺好，既用到了二叉树，还考了图中的建图和广度优先搜索的内容。涉及到的c++内容比如哈希表、集合的使用也比较多。这个题目出的确实有水平。

代码

#include<iostream>
using namespace std;

#include<unordered_map>
#include<unordered_set>
#include<functional>
#include<queue>

//Definition for a binary tree node.
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:
    int amountOfTime(TreeNode* root, int start) 
    {
        unordered_map<int, vector<int>> graph;
        function<void(TreeNode*)> dfs = [&](TreeNode* node)
            {
                for (TreeNode* child : vector<TreeNode*>{ node->left, node->right })
                {
                    if (child != nullptr)
                    {
                        graph[node->val].push_back(child->val);
                        graph[child->val].push_back(node->val);
                        dfs(child);
                    }
                }
            };
        dfs(root); // 哈希表建邻接图表
       
        // 广度优先搜索
        queue<vector<int>> q;
        q.push({ start, 0 }); // 该位置的值, 距离
        unordered_set<int> visited;
        visited.insert(start);
        int time = 0;
        while (q.empty() == false)
        {
            auto arr = q.front();
            q.pop();
            int nodeVal = arr[0];
            time = arr[1];
            for (int childVal : graph[nodeVal])
            {
                if (visited.count(childVal) == 0)
                {
                    q.push({ childVal, time + 1 });
                    visited.insert(childVal);
                }
            }
        }
        return time;
    }
};

int main()
{
    Solution s;


    return 0;
}