[Leetcode] Path with Maximum Probability - java

https://leetcode.com/problems/path-with-maximum-probability

 

문제 파악

무방향 그래프가 주어지고, 각 간선에는 가중치가 확률로 주어져있다.

시작 노드와 끝 노드가 주어지고, 끝 지점에 도달하는 최대 확률을 구하는 문제

 

입출력 예시

• Input: n = 3, edges = [[0,1],[1,2],[0,2]], succProb = [0.5,0.5,0.2], start = 0, end = 2
• Output: 0.25000

0.5 * 0.5 이기 때문에 최대 확률은 0.25가 된다.

---

접근 방법

가중치 그래프이기 때문에 다익스트라 알고리즘을 사용한다.

최소 비용이 아닌 최대 확률이기 때문에 우선순위는 큰 수부터 나오도록 설정해야 한다.

 

구현 순서

1. 인접리스트 초기화
2. 무방향이므로 둘 다 연결 후 확률을 넣어줌
3. 확률 초기화는 1로 진행 곱해야 하기 때문에
4. 큰 값을 우선순위로 둠. 최소 비용이 아닌 최대 확률을 구하는 문제이기 때문에
5. 그래서 확률이 더 작을 때 갱신

---

코드 구현

class Solution {
    public class Edge{
        int node;
        double prob;

        public Edge(int node, double prob){
            this.node = node;
            this.prob = prob;
        }
    }
    public double maxProbability(int n, int[][] edges, double[] succProb, int start_node, int end_node) {
        // 인접리스트 초기화
        List<List<Edge>> graph = new ArrayList<>();

        for(int i = 0; i < n; i++){
            graph.add(new ArrayList<>());
        }
        int i = 0;
        for(int[] edge : edges){
            int u = edge[0];
            int v = edge[1];
            double p = succProb[i];

            graph.get(u).add(new Edge(v, p));
            graph.get(v).add(new Edge(u, p));
            i++;
        }

        // 다익스트라
        double[] probability = new double[n];
        Arrays.fill(probability, 0.);

        Queue<Edge> pq = new PriorityQueue<>(Comparator.comparingDouble((Edge e) -> e.prob).reversed());

        pq.offer(new Edge(start_node, 1.)); // 확률은 곱해야 하니까 1.0로 초기화
        probability[start_node] = 1.;

        while(!pq.isEmpty()){
            // 방문
            Edge cur = pq.poll();

            // 도착하면 종료
            if(cur.node == end_node) return probability[end_node];

            // 예약
            for(Edge next : graph.get(cur.node)){
                double nextProb = probability[cur.node] * next.prob;
                if(probability[next.node] < nextProb){
                    pq.offer(new Edge(next.node, nextProb));
                    probability[next.node] = nextProb;
                }
            }
        }

        return 0.;
    }
}