Path with Maximum Probability: Modified Dijkstra's Algorithm

You have a network of nodes connected by edges, and each edge has a success probability. You want to find the path from a start node to an end node that gives you the highest chance of success. That is LeetCode 1514: Path with Maximum Probability, and it is one of the best problems for learning how to adapt Dijkstra's algorithm to work with probabilities instead of distances.

The problem

You are given an undirected weighted graph with n nodes (labeled 0 to n - 1), a list of edges, and a list of success probabilities for each edge. Given a start node and an end node, find the path with the maximum probability of success. If there is no path from start to end, return 0.

Probabilities are between 0 and 1, and the probability of a path is the product of the probabilities along its edges. For example, if a path goes through edges with probabilities 0.5 and 0.5, the total probability is 0.5 * 0.5 = 0.25.

Source nodeProcessed nodeTarget node

The approach

This is a single-source shortest path problem with a twist. Instead of minimizing distance (where you add edge weights), you are maximizing probability (where you multiply edge weights). The structure is identical to Dijkstra's algorithm, but with two key changes:

1. Use a max-heap instead of a min-heap. You want the node with the highest probability first, not the lowest distance.
2. Multiply instead of add. When you relax an edge, the new probability is prob[current] * edge_weight, not dist[current] + edge_weight.

Here is the plan:

1. Build an adjacency list from the edge list, storing both the neighbor and the edge probability
2. Initialize probabilities: every node starts at 0 except the source, which starts at 1.0 (certainty)
3. Use a max-heap: push (-1.0, start) onto the heap (negate because Python's heapq is a min-heap)
4. Process each node greedily: pop the node with the highest probability. If it is the target, return the probability. Otherwise, push all neighbors whose probability would improve
5. Return 0 if the target is never reached

The key insight is the same as classic Dijkstra's: once you pop a node from the heap, its probability is final. Since all probabilities are between 0 and 1, multiplying can only decrease the value. A longer path through a farther node cannot produce a higher probability than what you already have.

import heapq
from collections import defaultdict

def max_probability(n: int, edges: list[list[int]], succ_prob: list[float],
                    start: int, end: int) -> float:
    graph = defaultdict(list)
    for i, (a, b) in enumerate(edges):
        graph[a].append((b, succ_prob[i]))
        graph[b].append((a, succ_prob[i]))

    prob = [0.0] * n
    prob[start] = 1.0
    heap = [(-1.0, start)]

    while heap:
        neg_p, node = heapq.heappop(heap)
        p = -neg_p
        if node == end:
            return p
        if p < prob[node]:
            continue
        for neighbor, edge_prob in graph[node]:
            new_prob = p * edge_prob
            if new_prob > prob[neighbor]:
                prob[neighbor] = new_prob
                heapq.heappush(heap, (-new_prob, neighbor))

    return 0.0

Notice a few differences from the standard Dijkstra's you might have seen in Network Delay Time. First, the prob array is initialized to 0 instead of infinity, because we are maximizing rather than minimizing. Second, the source starts at 1.0 (100% probability of being at the start) rather than 0. Third, we negate probabilities when pushing to the heap because Python only provides a min-heap, and we need max-heap behavior.

The skip condition p < prob[node] is the analog of node in dist from the classic version. If we have already found a better probability for this node, the current heap entry is stale and we skip it.

Step-by-step walkthrough

Let's trace through the algorithm on our example graph with edges [[0,1], [1,2], [0,2]], probabilities [0.5, 0.5, 0.2], start = 0, end = 2.

The max-heap ensures we always process the node with the highest known probability first. When node 1 is processed with probability 0.5, it offers node 2 a probability of 0.5 * 0.5 = 0.25, which beats the direct path from node 0 at 0.2. By the time we pop node 2, we have the best possible probability.

Complexity analysis

Time: Each edge is relaxed at most a constant number of times, and each relaxation involves a heap push costing O(log V). With E edges total, the overall time is O(E log V).

Space: The adjacency list stores all E edges (each twice for the undirected graph), taking O(V + E). The heap can hold at most O(E) entries. The probability array holds V entries. Total space is O(V + E).

The building blocks

Modified Dijkstra's with max-heap

Classic Dijkstra's minimizes a sum of weights. This variant maximizes a product of weights. The structural change is small, but understanding why it works requires understanding the greedy property: since all probabilities are in the range [0, 1], multiplying by any edge probability can only decrease the current value. That means once you pop a node from the max-heap, no future path through other nodes can improve its probability. This is exactly analogous to how non-negative edge weights guarantee Dijkstra's correctness for shortest paths.

This same "flip the operation" trick applies to other graph problems. If you need to find the path that maximizes the minimum edge weight (bottleneck shortest path), you replace addition with min and use a max-heap. The key question is always: does the greedy property still hold?

Adjacency list for undirected weighted graphs

Unlike the directed graph in Network Delay Time, this graph is undirected. That means each edge gets added in both directions:

graph = defaultdict(list)
for i, (a, b) in enumerate(edges):
    graph[a].append((b, succ_prob[i]))
    graph[b].append((a, succ_prob[i]))

The rest of the algorithm does not change. The skip condition (p < prob[node]) handles the fact that an undirected edge might try to "go back" to an already-processed node. That entry gets popped from the heap and immediately skipped.

Edge cases

• No path exists: If start and end are in different connected components, the heap empties without ever reaching end. The function returns 0.0.
• Start equals end: The probability is 1.0. The algorithm handles this naturally: we push (-1.0, start), pop it, and since node == end, we return 1.0.
• All probabilities are 1.0: Every path has probability 1.0. The algorithm returns 1.0 as soon as it pops the end node.
• Single edge: If there is only one edge from start to end, the answer is that edge's probability. The algorithm pushes the neighbor, pops it, and returns.
• Multiple paths with the same probability: The heap may contain duplicate entries for the same node. The skip condition ensures we process each node at most once with its best probability.

From understanding to recall

You have seen how Dijkstra's algorithm adapts to probability maximization. The changes are minimal: swap the min-heap for a max-heap, multiply instead of add, initialize to 0 instead of infinity, and start the source at 1.0 instead of 0. The logic makes sense now. But reproducing it under interview pressure is a different challenge.

The details trip people up. Negating probabilities for Python's min-heap. Using p < prob[node] instead of node in dist. Returning early when you pop the target. These are small things, but forgetting any one of them produces a wrong answer.

Spaced repetition makes these details automatic. You practice writing the max-heap Dijkstra's loop, the probability relaxation, and the early return from scratch at increasing intervals. After a few rounds, the pattern becomes second nature. When a graph optimization problem appears in your interview, you do not need to re-derive the modifications. You just write them.

Related posts

• Network Delay Time - Classic Dijkstra's algorithm, the foundation this problem builds on
• Cheapest Flights Within K Stops - Another shortest path variant with additional constraints
• Course Schedule - Graph traversal with topological ordering