Skip to content

2385. Amount of Time for Binary Tree to Be Infected 👍

  • Time: $O(n)$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
class Solution {
 public:
  int amountOfTime(TreeNode* root, int start) {
    int ans = -1;
    const unordered_map<int, vector<int>> graph = getGraph(root);
    queue<int> q{{start}};
    unordered_set<int> seen{start};

    for (; !q.empty(); ++ans) {
      for (int sz = q.size(); sz > 0; --sz) {
        const int u = q.front();
        q.pop();
        if (!graph.count(u))
          continue;
        for (const int v : graph.at(u)) {
          if (seen.count(v))
            continue;
          q.push(v);
          seen.insert(v);
        }
      }
    }

    return ans;
  }

 private:
  unordered_map<int, vector<int>> getGraph(TreeNode* root) {
    unordered_map<int, vector<int>> graph;
    queue<pair<TreeNode*, int>> q{{{root, -1}}};  // (node, parent)

    while (!q.empty()) {
      const auto [node, parent] = q.front();
      q.pop();
      if (parent != -1) {
        graph[parent].push_back(node->val);
        graph[node->val].push_back(parent);
      }
      if (node->left)
        q.emplace(node->left, node->val);
      if (node->right)
        q.emplace(node->right, node->val);
    }

    return graph;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
class Solution {
  public int amountOfTime(TreeNode root, int start) {
    int ans = -1;
    Map<Integer, List<Integer>> graph = getGraph(root);
    Queue<Integer> q = new ArrayDeque<>(Arrays.asList(start));
    Set<Integer> seen = new HashSet<>(Arrays.asList(start));

    for (; !q.isEmpty(); ++ans) {
      for (int sz = q.size(); sz > 0; --sz) {
        final int u = q.poll();
        if (!graph.containsKey(u))
          continue;
        for (final int v : graph.get(u)) {
          if (seen.contains(v))
            continue;
          q.offer(v);
          seen.add(v);
        }
      }
    }

    return ans;
  }

  private Map<Integer, List<Integer>> getGraph(TreeNode root) {
    Map<Integer, List<Integer>> graph = new HashMap<>();
    // (node, parent)
    Queue<Pair<TreeNode, Integer>> q = new ArrayDeque<>(Arrays.asList(new Pair<>(root, -1)));

    while (!q.isEmpty()) {
      Pair<TreeNode, Integer> pair = q.poll();
      TreeNode node = pair.getKey();
      final int parent = pair.getValue();
      if (parent != -1) {
        graph.putIfAbsent(parent, new ArrayList<>());
        graph.putIfAbsent(node.val, new ArrayList<>());
        graph.get(parent).add(node.val);
        graph.get(node.val).add(parent);
      }
      if (node.left != null)
        q.add(new Pair<>(node.left, node.val));
      if (node.right != null)
        q.add(new Pair<>(node.right, node.val));
    }

    return graph;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
class Solution:
  def amountOfTime(self, root: Optional[TreeNode], start: int) -> int:
    ans = -1
    graph = self._getGraph(root)
    q = collections.deque([start])
    seen = {start}

    while q:
      ans += 1
      for _ in range(len(q)):
        u = q.popleft()
        if u not in graph:
          continue
        for v in graph[u]:
          if v in seen:
            continue
          q.append(v)
          seen.add(v)

    return ans

  def _getGraph(self, root: Optional[TreeNode]) -> Dict[int, List[int]]:
    graph = collections.defaultdict(list)
    q = collections.deque([(root, -1)])  # (node, parent)

    while q:
      node, parent = q.popleft()
      if parent != -1:
        graph[parent].append(node.val)
        graph[node.val].append(parent)
      if node.left:
        q.append((node.left, node.val))
      if node.right:
        q.append((node.right, node.val))

    return graph