Skip to content

1514. Path with Maximum Probability 👍

  • Time: $O(|E|\log |E|)$
  • Space: $O(|V| + |E|)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
class Solution {
 public:
  double maxProbability(int n, vector<vector<int>>& edges,
                        vector<double>& succProb, int start, int end) {
    // {a: [(b, probability_ab)]}
    vector<vector<pair<int, double>>> graph(n);
    // (the probability to reach u, u)
    priority_queue<pair<double, int>> maxHeap;
    maxHeap.emplace(1.0, start);
    vector<bool> seen(n);

    for (int i = 0; i < edges.size(); ++i) {
      const int u = edges[i][0];
      const int v = edges[i][1];
      const double prob = succProb[i];
      graph[u].emplace_back(v, prob);
      graph[v].emplace_back(u, prob);
    }

    while (!maxHeap.empty()) {
      const auto [prob, u] = maxHeap.top();
      maxHeap.pop();
      if (u == end)
        return prob;
      if (seen[u])
        continue;
      seen[u] = true;
      for (const auto& [nextNode, edgeProb] : graph[u]) {
        if (seen[nextNode])
          continue;
        maxHeap.emplace(prob * edgeProb, nextNode);
      }
    }

    return 0;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
class Solution {
  public double maxProbability(int n, int[][] edges, double[] succProb, int start, int end) {
    // {a: [(b, probability_ab)]}
    List<Pair<Integer, Double>>[] graph = new List[n];
    // (the probability to reach u, u)
    Queue<Pair<Double, Integer>> maxHeap =
        new PriorityQueue<>((a, b) -> Double.compare(b.getKey(), a.getKey())) {
          { offer(new Pair<>(1.0, start)); }
        };
    boolean[] seen = new boolean[n];

    for (int i = 0; i < n; ++i)
      graph[i] = new ArrayList<>();

    for (int i = 0; i < edges.length; ++i) {
      final int u = edges[i][0];
      final int v = edges[i][1];
      final double prob = succProb[i];
      graph[u].add(new Pair<>(v, prob));
      graph[v].add(new Pair<>(u, prob));
    }

    while (!maxHeap.isEmpty()) {
      final double prob = maxHeap.peek().getKey();
      final int u = maxHeap.poll().getValue();
      if (u == end)
        return prob;
      if (seen[u])
        continue;
      seen[u] = true;
      for (Pair<Integer, Double> node : graph[u]) {
        final int nextNode = node.getKey();
        final double edgeProb = node.getValue();
        if (seen[nextNode])
          continue;
        maxHeap.add(new Pair<>(prob * edgeProb, nextNode));
      }
    }

    return 0;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
class Solution:
  def maxProbability(
      self,
      n: int,
      edges: list[list[int]],
      succProb: list[float],
      start: int,
      end: int,
  ) -> float:
    graph = [[] for _ in range(n)]  # {a: [(b, probability_ab)]}
    maxHeap = [(-1.0, start)]   # (the probability to reach u, u)
    seen = [False] * n

    for i, ((u, v), prob) in enumerate(zip(edges, succProb)):
      graph[u].append((v, prob))
      graph[v].append((u, prob))

    while maxHeap:
      prob, u = heapq.heappop(maxHeap)
      prob *= -1
      if u == end:
        return prob
      if seen[u]:
        continue
      seen[u] = True
      for nextNode, edgeProb in graph[u]:
        if seen[nextNode]:
          continue
        heapq.heappush(maxHeap, (-prob * edgeProb, nextNode))

    return 0