LeetCode/java/path-with-maximum-probability.java

import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.PriorityQueue;

class Solution {
  private record Edge(int v, double p) {}

  private static Map<Integer, List<Edge>> makeGraph(int[][] edges, double[] p) {
    Map<Integer, List<Edge>> graph = new HashMap<>();

    for (int i = 0; i < edges.length; ++i) {
      var u = edges[i][0];
      var v = edges[i][1];
      var prob = p[i];

      graph.computeIfAbsent(u, k -> new ArrayList<>()).add(new Edge(v, prob));
      graph.computeIfAbsent(v, k -> new ArrayList<>()).add(new Edge(u, prob));
    }

    return graph;
  }

  private record Entry(double probability, int u) {}

  public double maxProbability(
      int n, int[][] edges, double[] succProb, int startNode, int endNode) {
    var graph = makeGraph(edges, succProb);
    var maxProb = new double[n];

    PriorityQueue<Entry> q =
        new PriorityQueue<>(Comparator.comparingDouble(Entry::probability).reversed());

    maxProb[startNode] = 1;
    q.add(new Entry(maxProb[startNode], startNode));

    while (!q.isEmpty()) {
      var candidate = q.poll();
      if (candidate.u() == endNode) {
        return candidate.probability();
      }

      for (var edge : graph.getOrDefault(candidate.u(), new ArrayList<>())) {
        var possibleProbability = candidate.probability() * edge.p();

        if (possibleProbability > maxProb[edge.v()]) {
          maxProb[edge.v()] = possibleProbability;
          q.add(new Entry(maxProb[edge.v()], edge.v()));
        }
      }
    }

    return 0;
  }
}