Update frog-position-after-t-seconds.py

Author: kamyu104

Python/frog-position-after-t-seconds.py

@@ -5,6 +5,37 @@
 
 
 class Solution(object):
+    def frogPosition(self, n, edges, t, target):
+        """
+        :type n: int
+        :type edges: List[List[int]]
+        :type t: int
+        :type target: int
+        :rtype: float
+        """        
+        def dfs(G, target, t, node, parent):
+            if not t or not (len(G[node])-(parent != 0)):
+                return int(node == target)
+            result = 0
+            for child in G[node]:
+                if child == parent:
+                    continue
+                result = dfs(G, target, t-1, child, node)
+                if result:
+                    break
+            return result*(len(G[node])-(parent != 0))
+        
+        G = collections.defaultdict(list)
+        for u, v in edges:
+            G[u].append(v)
+            G[v].append(u)
+        choices = dfs(G, target, t, 1, 0)
+        return 1.0/choices if choices else 0.0
+
+
+# Time:  O(n)
+# Space: O(n)
+class Solution2(object):
     def frogPosition(self, n, edges, t, target):
         """
         :type n: int
@@ -16,8 +47,12 @@ def frogPosition(self, n, edges, t, target):
         def dfs(G, target, t, node, parent):
             if not t or not (len(G[node])-(parent != 0)):
                 return float(node == target)
-            result = sum(dfs(G, target, t-1, child, node)
-                         for child in G[node] if child != parent)
+            for child in G[node]:
+                if child == parent:
+                    continue
+                result = dfs(G, target, t-1, child, node)
+                if result:
+                    break
             return result/(len(G[node])-(parent != 0))
 
         G = collections.defaultdict(list)