Path with Maximum Probability Python Solution

Tested Python solution for LeetCode 1514 with 16 pytest cases. Generate a practice environment with lcpy.

LeetCode 1514, Medium. Topics: Array, Graph Theory, Heap (Priority Queue), Shortest Path, Dijkstra’s Algorithm. View on LeetCode.

Generate this problem as a practice environment: tested reference solution, 16 parametrized pytest cases, and a playground notebook:

lcpy gen -n 1514   # by problem number
lcpy gen -s path_with_maximum_probability   # by problem name

​

Problem

You are given an undirected weighted graph of n nodes (0-indexed), represented by an edge list where edges[i] = [a<sub>i</sub>, b<sub>i</sub>] is an undirected edge connecting the nodes a<sub>i</sub> and b<sub>i</sub> 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 end, return 0. Your answer will be accepted if it differs from the correct answer by at most 10^-5.

​

Examples

Input: n = 3, edges = [[0,1],[1,2],[0,2]], succProb = [0.5,0.5,0.2], start = 0, end = 2
Output: 0.25000
Explanation: There are two paths from start to end, one having a probability of success = 0.2 and the other has 0.5 * 0.5 = 0.25.

Input: n = 3, edges = [[0,1],[1,2],[0,2]], succProb = [0.5,0.5,0.3], start = 0, end = 2
Output: 0.30000

Input: n = 3, edges = [[0,1]], succProb = [0.5], start = 0, end = 2
Output: 0.00000
Explanation: There is no path between 0 and 2.

​

Constraints

• 2 <= n <= 10^4
• 0 <= start, end < n
• start != end
• 0 <= a_i, b_i < n
• a_i != b_i
• 0 <= succProb.length == edges.length <= 2 * 10^4
• 0 <= succProb[i] <= 1
• There is at most one edge between every two nodes.

​

Solution

Reference implementation from solution.py on GitHub, full suite in test_solution.py:

import heapq


class Solution:
    # Time: O((V + E) * log V)
    # Space: O(V + E)
    def max_probability(
        self, n: int, edges: list[list[int]], succ_prob: list[float], start_node: int, end_node: int
    ) -> float:
        adj: list[list[tuple[int, float]]] = [[] for _ in range(n)]
        for (a, b), p in zip(edges, succ_prob, strict=True):
            adj[a].append((b, p))
            adj[b].append((a, p))

        best = [0.0] * n
        best[start_node] = 1.0
        heap: list[tuple[float, int]] = [(-1.0, start_node)]
        while heap:
            neg, node = heapq.heappop(heap)
            cur = -neg
            if cur < best[node]:
                continue
            if node == end_node:
                return cur
            for nxt, p in adj[node]:
                cand = cur * p
                if cand > best[nxt]:
                    best[nxt] = cand
                    heapq.heappush(heap, (-cand, nxt))
        return best[end_node]

​

Complexity

Time	Space
O((V + E) * log V)	O(V + E)

​

Tags

NeetCode All.