Problem statement

https://leetcode.com/problems/frog-position-after-t-seconds/

Solution

Nothing very interesting here, we need to traverse our tree and keep number of steps and probability. If we reached target, we need to check:

  1. if q == 1 and par != 0, then we have leaf and we need condition steps <= t.
  2. if q > 1 then we not in the leaf and we need steps == t.
  3. if q = 0, it only can happen for empty tree.

If we reached leaf, return. Then for each of children traverse them, using prob/(q-1+(par==0)) condition: if we were in the root we divide by q, if not: by q - 1. Finally, we run dfs(0, 1, 0, 1), where first argument 0 is non-existing parent.

Complexity

Time complexity is O(n), space is O(H).

Code

class Solution:
    def frogPosition(self, n, edges, t, target):
        graph = defaultdict(list)
        for x, y in edges:
            graph[x].append(y)
            graph[y].append(x)
            
        self.ans = None    

        def dfs(par, node, steps, prob):
            q = len(graph[node])
            if node == target:
                if (q == 1 and steps <= t and par != 0) or (q > 1 and steps == t) or (q == 0):
                    self.ans = prob
                else:
                    self.ans = 0
                   
            if q == 1 and par != 0:
                return
            for neib in graph[node]:
                if neib == par: continue
                dfs(node, neib, steps+1, prob/(q-1+(par==0)))
                
        dfs(0, 1, 0, 1)
        return self.ans