Skip to content

1377. Frog Position After T Seconds 👍

  • 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
class Solution {
 public:
  double frogPosition(int n, vector<vector<int>>& edges, int t, int target) {
    vector<vector<int>> tree(n + 1);
    queue<int> q{{1}};
    vector<bool> seen(n + 1);
    vector<double> prob(n + 1);

    seen[1] = true;
    prob[1] = 1.0;

    for (const vector<int>& edge : edges) {
      const int u = edge[0];
      const int v = edge[1];
      tree[u].push_back(v);
      tree[v].push_back(u);
    }

    while (!q.empty() && t-- > 0)
      for (int sz = q.size(); sz > 0; --sz) {
        const int a = q.front();
        q.pop();
        const int nChildren = count_if(tree[a].begin(), tree[a].end(),
                                       [&seen](int b) { return !seen[b]; });
        for (const int b : tree[a]) {
          if (seen[b])
            continue;
          seen[b] = true;
          prob[b] = prob[a] / nChildren;
          q.push(b);
        }
        if (nChildren > 0)
          prob[a] = 0.0;
      }

    return prob[target];
  }
};

 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
class Solution {
  public double frogPosition(int n, int[][] edges, int t, int target) {
    List<Integer>[] tree = new List[n + 1];
    Queue<Integer> q = new ArrayDeque<>(Arrays.asList(1));
    boolean[] seen = new boolean[n + 1];
    double[] prob = new double[n + 1];

    seen[1] = true;
    prob[1] = 1.0;

    for (int i = 1; i <= n; ++i)
      tree[i] = new ArrayList<>();

    for (int[] edge : edges) {
      final int u = edge[0];
      final int v = edge[1];
      tree[u].add(v);
      tree[v].add(u);
    }

    while (!q.isEmpty() && t-- > 0)
      for (int sz = q.size(); sz > 0; --sz) {
        final int a = q.poll();
        int nChildren = 0;
        for (final int b : tree[a])
          if (!seen[b])
            ++nChildren;
        for (final int b : tree[a]) {
          if (seen[b])
            continue;
          seen[b] = true;
          prob[b] = prob[a] / nChildren;
          q.add(b);
        }
        if (nChildren > 0)
          prob[a] = 0.0;
      }

    return prob[target];
  }
}

 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
class Solution:
  def frogPosition(self, n: int, edges: List[List[int]], t: int, target: int) -> float:
    tree = [[] for _ in range(n + 1)]
    q = collections.deque([1])
    seen = [False] * (n + 1)
    prob = [0] * (n + 1)

    prob[1] = 1
    seen[1] = True

    for u, v in edges:
      tree[u].append(v)
      tree[v].append(u)

    for _ in range(t):
      for _ in range(len(q)):
        a = q.popleft()
        nChildren = sum(not seen[b] for b in tree[a])
        for b in tree[a]:
          if seen[b]:
            continue
          seen[b] = True
          prob[b] = prob[a] / nChildren
          q.append(b)
        if nChildren > 0:
          prob[a] = 0

    return prob[target]