1514. Path with Maximum Probability

main.py

class Solution(object):
    def maxProbability(self, n, edges, succProb, start, end):

        # it looks like dijsktra but with different relaxation
        g = collections.defaultdict(dict)
        for i in range(len(edges)):
            s = edges[i][0]
            d = edges[i][1]
            p = succProb[i]
            g[s][d] = g[d][s] = p
                    
        dist = [0 for _ in range(n)]
        dist[start] = 1
        q = []
        heapq.heappush(q, (1, start))
        while q:
            _, u = heapq.heappop(q)
            for v, w in g[u].items():
                new_dist = dist[u] * w
                if new_dist > dist[v]:
                    dist[v] = new_dist
                    heapq.heappush(q, (-math.log(dist[u] * w), v))
        return dist[end]