题目链接:https://leetcode.cn/problems/amount-of-time-for-binary-tree-to-be-infected/
题目大意:给一棵二叉树,其中元素各异。给一个节点的元素start
,第0
分钟时它是感染源。然后之后的每一分钟,被感染的节点都会传染其儿子节点和父亲节点。求整棵树被感染的时间。
思路:DFS+BFS。DFS记录邻居节点(父亲节点和儿子节点),BFS模拟感染过程。
完整代码
/**
* 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:
vector<int> neb[100001];
vector<bool> infected;
set<int> uninf;
void dfs(TreeNode* node, int parent) {
if (node == nullptr)
return;
uninf.insert(node->val);
if (parent != 0)
neb[node->val].push_back(parent);
if (node->left) {
neb[node->val].push_back(node->left->val);
dfs(node->left, node->val);
}
if (node->right) {
neb[node->val].push_back(node->right->val);
dfs(node->right, node->val);
}
}
void infect(int infected, int mnt) {
}
int amountOfTime(TreeNode* root, int start) {
dfs(root, 0);
infected.resize(100001, false);
int mnt = -1;
vector<int> nxt;
vector<int> now_inf;
nxt.push_back(start);
while(!uninf.empty()) {
mnt++;
now_inf = nxt;
nxt.clear();
for (int i = 0; i < now_inf.size(); i++) {
infected[now_inf[i]] = true;
uninf.erase(now_inf[i]);
if (uninf.empty())
break;
for (int j = 0; j < neb[now_inf[i]].size(); j++) {
if (!infected[neb[now_inf[i]][j]]) {
nxt.push_back(neb[now_inf[i]][j]);
}
}
}
}
return mnt;
}
};