1514. Path with Maximum Probability

from collections import defaultdict
import heapq
class Solution:
    def maxProbability(self, n: int, edges: List[List[int]], succProb: List[float], start: int, end: int) -> float:
        '''
        Input
            n - number of nodes
            edges - (u,v) for an undirected edge.
            succProb - edge[i]'s weight
            start - start node
            end - end node
        Output
            MAXIMUM probability for a given path
        Constraint
            Taking multiple edges MULTIPLIES their weights
        Approach
            Dijkstra's algorithm to find Optimal path
            Optimal:
                Highest probability.
                Or use successProb directly and find MAXIMUM path
        
        '''
        adjList = defaultdict(list)
        
        for idx, (u, v) in enumerate(edges):
            weight = succProb[idx]
            
            adjList[u].append((v, weight))
            adjList[v].append((u, weight))
        
        q = []
        
        heapq.heappush(q, (-1, start))
        
        visited = [False] * n
        while q:
            cumCost, node = heapq.heappop(q)
            cumCost = -1 * cumCost
            
            if visited[node]:
                continue
            
            visited[node] = True
            # Early return works because this is a greedy approach
            if node == end:
                return cumCost
            
            # BFS
            for to, weight in adjList[node]:
                heapq.heappush(q, (-1 * cumCost * weight, to))
        return 0.0