LeetCode 1377. Frog Position After T Seconds

, , , , ,

Question

Given an undirected tree consisting of n vertices numbered from 1 to n. A frog starts jumping from the vertex 1. In one second, the frog jumps from its current vertex to another unvisited vertex if they are directly connected. The frog can not jump back to a visited vertex. In case the frog can jump to several vertices it jumps randomly to one of them with the same probability, otherwise, when the frog can not jump to any unvisited vertex it jumps forever on the same vertex.

The edges of the undirected tree are given in the array edges, where edges[i] = [fromi, toi] means that exists an edge connecting directly the vertices fromi and toi.

Return the probability that after t seconds the frog is on the vertex target.

Example

1
2
3
4
5
6
7
Input: n = 7, edges = [[1,2],[1,3],[1,7],[2,4],[2,6],[3,5]], t = 2, target = 4
Output: 0.16666666666666666 
Explanation: The figure above shows the given graph. The frog starts at vertex 1, jumping with 1/3 probability to the vertex 2 after second 1 and then jumping with 1/2 probability to vertex 4 after second 2. Thus the probability for the frog is on the vertex 4 after 2 seconds is 1/3 * 1/2 = 1/6 = 0.16666666666666666. 

Input: n = 7, edges = [[1,2],[1,3],[1,7],[2,4],[2,6],[3,5]], t = 1, target = 7
Output: 0.3333333333333333
Explanation: The figure above shows the given graph. The frog starts at vertex 1, jumping with 1/3 = 0.3333333333333333 probability to the vertex 7 after second 1.

Solution

首先理解题目意思，只要有选择可以跳，frog就会在t步内继续跳。这里也涉及到如果target是在叶子节点，然后无法再跳，我们就可以不管t的大小。但是如果target在途中的话，我们还有多的选择以及剩余t值，那么就无法到达。这道题也是一个无向图，所以构建的时候也需要小心。

Code

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
// time: O(n) space:O(n)
   public double frogPosition(int n, final int[][] edges, int t, final int target) {
       List<Integer>[] graph = new List[n];
       for (int i = 0; i < n; i++) graph[i] = new ArrayList<>();
       for (final int [] e : edges) {
           graph[e[0] - 1].add(e[1] - 1);
           graph[e[1] - 1].add(e[0] - 1);
       }
       final boolean[] visited = new boolean[n];
       visited[0] = true;
       final double[] prob = new double[n];
       prob[0]  = 1.0;
       final Queue<Integer> queue = new LinkedList<>();
       queue.offer(0);
       // 要考虑到frog不能跳，只有没有选择的时候，不能跳的时候才能留在原点跳完并且保留其概率
       while (!queue.isEmpty() && t-- > 0) {
           final int size = queue.size();
           for (int i = 0; i < size; i++) {
               final int curr = queue.poll();
               int adjCount = 0;
               for (final int adj : graph[curr]) if (!visited[adj]) adjCount++;
               for (final int adj : graph[curr]) {
                   if (!visited[adj]) {
                       visited[adj] = true;
                       queue.offer(adj);
                       prob[adj] = prob[curr] / adjCount;
                   }
               }
               // 很关键，我们只需要上面保留下一层子节点的概率值就好
               if (adjCount > 0) prob[curr] = 0;// 因为不是目标地方，还会继续走
           }
       }
       return prob[target - 1];
   }

Reference