[Leetcode]1514. Path with Maximum Probability

You are given an undirected weighted graph of n nodes (0-indexed), represented by an edge list where edges[i] = [a, b] is an undirected edge connecting the nodes a and b with a probability of success of traversing that edge succProb[i].
Given two nodes start and end, find the path with the maximum probability of success to go from start to end and return its success probability.
If there is no path from start to endreturn 0. 
Example 1:
Input: n = 3, edges = [[0,1],[1,2],[0,2]], succProb = [0.5,0.5,0.2], start = 0, end = 2 
Output: 0.25000 

This question first appeared in Leetcode Weekly Contest 197. I first tried DFS for a while but kept receiving TLE. Then I realized there are lots of unnecessary computations with DFS:

Given two paths a->...->d->...->e and b->...->d->...->e, DFS would compute d->…->e for twice, but if we found P(a->...->d) > P(b->...d->), then we can actually prune the second path.

Therefore Dijkstra’s shortest path algorithm would help. Details about the explanation and implementations Dijkstra’s algorithm are available here at: Shortest Path Algorithms I – Dijkstra.

class Solution:
    def maxProbability(self, n: int, edges: List[List[int]], succProb: List[float], start: int, end: int) -> float:
        # construct the graph
        graph = collections.defaultdict(list)
        prob = collections.defaultdict(dict)
        for i,(p,[a,b]) in enumerate(zip(succProb,edges)):
            prob[a][b] = prob[b][a] = p
        # apply dijkstra
        dis = {start:1}
        for i in range(n):
            dis[i] = 0
        visited = set([])
        # note that Python only supports min-heap
        # so some tricks here to get a max-heap
        pq = [(-1,start)]
        while pq:
            _p, node = heapq.heappop(pq)
            for child in graph[node]:
                if child not in visited:
                    if dis[child] < -1 * _p * prob[node][child]:
                        heapq.heappush(pq,(_p * prob[node][child],child))
                        dis[child] = -1 * _p * prob[node][child]

        return dis[end]
Notify of
Inline Feedbacks
View all comments