1753 - 최단경로

전형적인 다익스트라 문제로 정점, 간선의 수와 시작 정점 그리고 간선이 입력으로 주어지며, 이를 입력받아 시작 정점에서 각 정점으로의 최단 거리를 구하는 문제이다.

이 문제를 풀기 위해 다익스트라를 구현하면 되며 이 때 한 노드에서 다른 노드로 가는 간선이 여러개일 수 있다는 사실에 주의해야 한다(가중치가 다를 수 있기 때문).

🔗 문제 링크

코드

간선과 다익스트라 탐색 상태를 나타내기 위한 클래스를 선언한다.

static class Edge {
    public int to;
    public int weight;

    public Edge(int to, int weight) {
        this.to = to;
        this.weight = weight;
    }
}

static class State {
    public int node;
    public long dist;

    public State(int node, long dist) {
        this.node = node;
        this.dist = dist;
    }
}

정점 수 V, 간선 수 E, 시작 정점 K를 입력받는다. 각 정점별 최단 이동거리를 저장할 dist 배열을 선언하고 초기값을 모두 INF로 설정해둔다. 그런 다음 각 정점별로 이동할 수 있는 간선을 담아둘 graph 변수를 선언한다.

참고로 INF는 가능한 정점까지의 최단거리의 집합에서의 최대값보다 더 큰 수여야 한다. 이 문제에서 간선 수는 최대 30만이고, 간선의 가중치는 최대 10이므로 어떠한 최단 경로도 3000005($E * 10 + 5$)보다 거리가 클 수는 없다. 왜냐하면 최단 경로는 단순 경로이므로 같은 간선을 반복하지 않기 때문이다. 따라서 INF = 3000005로 설정해 두었다.

public static void main(String[] args) throws IOException {
    StringTokenizer st = new StringTokenizer(br.readLine());
    int V = Integer.parseInt(st.nextToken());
    int E = Integer.parseInt(st.nextToken());
    int K = Integer.parseInt(br.readLine());

    long[] dist = new long[V+5];
    Arrays.fill(dist, INF);
    
    @SuppressWarnings("unchecked")
    List<Edge>[] graph = new List[V+5];
    for (int i = 0; i < V+5; i++)
        graph[i] = new ArrayList<>();

간선 정보를 입력받아 그래프에 추가한다. 이 때 반드시 두 정점에 대해 여러 간선이 존재할 수 있도록 하여야 한다. 또한 방향 그래프이므로 대칭적으로 간선을 삽입하면 안 된다.

그런 다음 dist 속성 기준 최소 힙 구조를 가지는 우선순위 큐 pq를 생성하고 탐색을 시작하는 최초 정점인 K를 넣는다. 정점 K->K 거리는 당연히 0이므로 dist[k] = 0으로 설정한다.

    for (int i = 1; i <= E; i++) {
        st = new StringTokenizer(br.readLine());

        int from = Integer.parseInt(st.nextToken());
        int to = Integer.parseInt(st.nextToken());
        int weight = Integer.parseInt(st.nextToken());

        graph[from].add(new Edge(to, weight));
    }

    PriorityQueue<State> pq 
        = new PriorityQueue<>(Comparator.comparingLong(e->e.dist));

    pq.offer(new State(K, 0));
    dist[K] = 0;

반복문을 돌면서 각 정점별 최소 거리를 구한다. 우선순위가 높은 State를 꺼내고, 그 State가 유효한지 확인한다. 그런 다음 그 State에서의 정점을 기준으로 인접 정점을 탐색하며, 더 짧은 거리로 업데이트할 수 있는지 확인한다. 가능하다면 거리를 업데이트하고 우선순위 큐에 넣는다.

    while (!pq.isEmpty()) {
        State s = pq.poll();
        if (s.dist > dist[s.node]) continue;

        for (Edge adj : graph[s.node]) {
            if (s.dist + adj.weight < dist[adj.to]) {
                dist[adj.to] = s.dist + adj.weight;
                pq.offer(new State(adj.to, dist[adj.to]));
            }
        }
    }

다익스트라가 끝나면 모든 정점별 최단 거리를 출력한다. 정점의 거리가 INF 값과 같다면 시작 정점 K에서 도달할 수 없는 정점임을 의미하므로 INF를 문자열로 출력한다.

    for (int i = 1; i <= V; i++) {
        bw.write((dist[i] == INF ? "INF" : dist[i]) + "\n");
    }
    bw.flush();

}

전체 코드

import java.util.*;
import java.io.*;

class Main {
    static BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    static BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
    static final long INF = 3000005L;

    static class Edge {
        public int to;
        public int weight;

        public Edge(int to, int weight) {
            this.to = to;
            this.weight = weight;
        }
    }

    static class State {
        public int node;
        public long dist;

        public State(int node, long dist) {
            this.node = node;
            this.dist = dist;
        }
    }

    public static void main(String[] args) throws IOException {
        StringTokenizer st = new StringTokenizer(br.readLine());
        int V = Integer.parseInt(st.nextToken());
        int E = Integer.parseInt(st.nextToken());
        int K = Integer.parseInt(br.readLine());

        long[] dist = new long[V+5];
        Arrays.fill(dist, INF);
        
        @SuppressWarnings("unchecked")
        List<Edge>[] graph = new List[V+5];
        for (int i = 0; i < V+5; i++)
            graph[i] = new ArrayList<>();


        for (int i = 1; i <= E; i++) {
            st = new StringTokenizer(br.readLine());

            int from = Integer.parseInt(st.nextToken());
            int to = Integer.parseInt(st.nextToken());
            int weight = Integer.parseInt(st.nextToken());

            graph[from].add(new Edge(to, weight));
        }

        PriorityQueue<State> pq 
            = new PriorityQueue<>(Comparator.comparingLong(e->e.dist));

        pq.offer(new State(K, 0));
        dist[K] = 0;

        while (!pq.isEmpty()) {
            State s = pq.poll();
            if (s.dist > dist[s.node]) continue;

            for (Edge adj : graph[s.node]) {
                if (s.dist + adj.weight < dist[adj.to]) {
                    dist[adj.to] = s.dist + adj.weight;
                    pq.offer(new State(adj.to, dist[adj.to]));
                }
            }
        }
        
        for (int i = 1; i <= V; i++) {
            bw.write((dist[i] == INF ? "INF" : dist[i]) + "\n");
        }
        bw.flush();

    }
}

시간복잡도

다익스트라 알고리즘의 시간복잡도는 $O((V+E)\log V)$ 로 알려져 있으므로 이 코드의 시간복잡도는 $O((V+E) \log V)$ 이다. 최악의 경우에도 $(20000 + 300000) \cdot \log_2 20000 \approx 320000 \cdot 14.3 = 4576000 \lt 1\text{억}$ 이므로 시간초과 없이 통과할 수 있다.