Problem statement

https://binarysearch.com/problems/Path-to-Maximize-Probability-to-Destination/

Solution

Equal to Leetcode 1514 Path with Maximum Probability but with more edge cases.

Complexity

Time complexity O(E + V * log V), space is O(E).

Code

class Solution:
    def solve(self, edges, success, start, end):
        n = 0
        G = defaultdict(list)
        
        for i, (a, b) in enumerate(edges):
            if success[i] == 0: continue
            G[a].append((b, -math.log(success[i])))
            G[b].append((a, -math.log(success[i])))
            n = max(n, a + 1, b + 1)
        
        heap = [(0, start)]
        dist = [float('inf')] * n
        if start >= n: return 0
        dist[start] = 0
        
        while heap:
            min_dist, idx = heappop(heap)
            if idx == end: return math.exp(-min_dist)
            for neibh, weight in G[idx]:
                cand = min_dist + weight
                if cand < dist[neibh]:
                    dist[neibh] = cand
                    heappush(heap, (cand, neibh))
                    
        return 0