leetcode

0 class Solution {

1 public:

2 int amountOfTime(TreeNode *root, int start) const {

3 unordered_map<TreeNode *, TreeNode *> parent;

4 queue<TreeNode *> q;

5 TreeNode *infected;

6

7 q.push(root);

8 while (!q.empty()) {

9 TreeNode *root = q.front();

10 q.pop();

11 if (root->val == start) {

12 infected = root;

13 break;

14 }

15

16 if (root->left) {

17 parent.insert({root->left, root});

18 q.push(root->left);

19 }

20

21 if (root->right) {

22 parent.insert({root->right, root});

23 q.push(root->right);

24 }

25 }

26

27 int depth = -1;

28 q = queue<TreeNode *>();

29 q.push(infected);

30 while (!q.empty()) {

31 depth++;

32 for (int k = q.size(); k > 0; k--) {

33 TreeNode *root = q.front();

34 q.pop();

35 if (parent.count(root)) {

36 TreeNode *prnt = parent[root];

37 (prnt->left == root ? prnt->left : prnt->right) = nullptr;

38 q.push(parent[root]);

39 }

40 if (root->left) {

41 q.push(root->left);

42 parent.erase(root->left);

43 }

44 if (root->right) {

45 q.push(root->right);

46 parent.erase(root->right);

47 }

48 }

49 }

50 return depth;

51 }

52 };