백준 23793 - 두 단계 최단 경로 1

문제

https://www.acmicpc.net/problem/23793

리뷰

다익스트라로 풀이할 수 있는 간단한 문제였다.

경로 $X\rightarrow Y \rightarrow Z$의 최단 거리는 $X$를 시작점으로
다익스트라를 수행하여 구한 dist[y]의 값과 $Y$를 시작점으로 다익스트라를
수행하여 구한 dist[z]의 값을 합하여 구할 수 있다.

마찬가지로, $X \rightarrow Z$의 경로는 다익스트라를 수행하여 구할 수 있는데
이 때 $Y$를 거치지 않아야 하므로 다익스트라 로직내에서 다음 이동할 정점이
$Y$인 간선들을 마주할 경우 넘어가도록 구성하였다.

로직의 시간복잡도는 다익스트라의 복잡도인 $O(NlogM)$로 수렴하고,
$N=100,000$, $M=200,000$인 최악의 경우에도 제한 조건 0.3초를 무난히
통과한다.

코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.*;

import static java.lang.Integer.*;

public class Main {
    static int[] dist;
    static List<List<Node>> graph = new ArrayList<>();

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        int N = parseInt(st.nextToken());
        int M = parseInt(st.nextToken());

        dist = new int[N + 1];
        for (int i = 0; i <= N; i++)
            graph.add(new ArrayList<>());

        int u, v, w;
        while (M-- > 0) {
            st = new StringTokenizer(br.readLine());
            u = parseInt(st.nextToken());
            v = parseInt(st.nextToken());
            w = parseInt(st.nextToken());

            graph.get(u).add(new Node(v, w));
        }

        int x, y, z;
        st = new StringTokenizer(br.readLine());
        x = parseInt(st.nextToken());
        y = parseInt(st.nextToken());
        z = parseInt(st.nextToken());

        int result1, result2;
        dijkstra(x, -1);
        result1 = dist[y] == MAX_VALUE ? -1 : dist[y];
        dijkstra(y, -1);
        if (result1 == -1 || dist[z] == MAX_VALUE)
            result1 = -1;
        else
            result1 += dist[z];

        dijkstra(x, y);
        result2 = dist[z] == MAX_VALUE ? -1 : dist[z];

        System.out.println(result1 + " " + result2);
        br.close();
    }

    static void dijkstra(int start, int y) {
        Arrays.fill(dist, MAX_VALUE);
        PriorityQueue<Node> pq = new PriorityQueue<>(Comparator.comparingInt(n -> n.w));
        dist[start] = 0;
        pq.offer(new Node(start, dist[start]));

        while (!pq.isEmpty()) {
            Node cur = pq.poll();

            if (dist[cur.v] < cur.w)
                continue;

            for (Node next : graph.get(cur.v)) {
                if (next.v == y) continue;

                if (dist[next.v] < dist[cur.v] + next.w)
                    continue;

                dist[next.v] = dist[cur.v] + next.w;
                pq.offer(new Node(next.v, dist[next.v]));
            }
        }
    }

    static class Node {
        int v, w;

        public Node(int v, int w) {
            this.v = v;
            this.w = w;
        }
    }
}

결과