Frog Position After T Seconds LeetCode solution in Python

1
class Solution:
2
    def frogPosition(
3
        self, n: int, edges: List[List[int]], t: int, target: int
4
    ) -> float:
5
        nei = collections.defaultdict(set)
6
        for a, b in edges:
7
            nei[a].add(b)
8
            nei[b].add(a)
9

10
        visited, res = set(), 0.0
11

12
        def dfs(leaf_id, p, time):
13
            nonlocal res
14
            if time >= t:
15
                if leaf_id == target:
16
                    res = p
17
                return
18
            visited.add(leaf_id)
19
            neighbors = nei[leaf_id] - visited
20
            for n in neighbors or [leaf_id]:
21
                dfs(n, p / (len(neighbors) or 1), time + 1)
22

23
        dfs(1, 1, 0)
24
        return res