Leetcode competition 179 weekly competitions (4) ------------ 1377. Frog position after T seconds

Give you an undirected tree consisting of n vertices, with vertices numbered from 1 to  n. The frog jumps from  vertex 1  . The rules are as follows:

• In the second, from which it is currently frog jump to another vertex  unvisited  vertex treated (if they are directly linked).
• The frog cannot jump back to the vertices that have been visited.
• If a frog can jump to many different vertices, then it has the same chance to jump to any one of them.
• If the frog cannot jump to any unvisited vertex, it will stay in place every time it jumps.

Tree edges with undirected array  edges is described, wherein  edges[i] = [fromi, toi] the presence of a direct communication means  fromi and the  toi edges of the two vertices.

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

 

Example 1:

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 jumping path of the frog. Frog 1 off from the apex, the first one second with a probability of 1/3 skip vertex 2, and the second 2 seconds with a probability of 1/2 skip vertex 4, so after 2 seconds frog positioned vertex probability is 1/3 4 * 1/2 = 1/6 = 0.16666666666666666. 

Example 2:

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 jumping path of the frog. The frog jumps from vertex 1, and has a probability of 1/3 = 0.3333333333333333 to jump to vertex 7 in 1 second . 

Example 3:

Input: n = 7, edges = [[1,2], [1,3], [1,7], [2,4], [2,6], [3,5]], t = 20, target = 6
 output: 0.16666666666666666

 

prompt:

• 1 <= n <= 100
• edges.length == n-1
• edges[i].length == 2
• 1 <= edges[i][0], edges[i][1] <= n
• 1 <= t <= 50
• 1 <= target <= n
• 10^-5 The result within the error of the accurate value  will be judged as correct.

 

My code:

class Solution:
    def frogPosition(self, n: int, edges: List[List[int]], t: int, target: int) -> float:
        from collections import defaultdict
        graph = defaultdict(set)
        for e in edges:
            start, end = e[0], e[1]
            if start > end:
                start, end = end, start
            graph[start].add(end)
            
        chance, s, step = 1, [(1,1)], 0
        while s and step <= t:
            n_s = []
            for node, chance in s:
                if node == target and step == t:
                    return chance
                
                if len(graph[node]):
                    n_chance = 1/len(graph[node])
                    for n in graph[node]:
                        n_s.append((n,chance*n_chance))
                else:
                    n_s.append((node,chance))
            step+=1
            s = n_s
        return 0