[실전 문제] 파티

1. 문제

철수네 마을에는 N개의 집이 있으며, 각 집은 고유한 번호를 부여받는다. 각 번호는 1보다 크거나 같고, N보다 작거나 같다. 철수는 마을 K에 살고 있다. 집과 집 사이에는 단방향 도로가 존재하기 때문에, 이 도로를 통하여 서로 이동할 수 있다. 즉, N개의 마을은 그래프 구조를 이룬다. 편의상 각 집에는 한 사람만이 살고 있다고 가정하자. 크리스마스인 오늘, 철수는 본인의 집에서 파티를 열려고 한다. 따라서 다른 모든 사람들이 철수의 집에 모여 파티를 즐기고, 파티가 끝난 후에는 다시 본인의 집으로 돌아가려 한다. 사람들은 본인의 집에서 철수네 집까지 이동하기 위하여 항상 최단거리로 이동하기를 원하고, 마찬가지로 철수네 집에서 본인의 집에 갈 때도 최단거리로 이동하기를 원한다. 집의 개수와 두 집을 잇는 단방향 도로의 정보, 그리고 철수의 집 번호가 주어질 때, 마을 사람들이 파티를 즐기기 위해서 이동해야 하는 총 거리의 합을 출력하는 프로그램을 작성하시오.

---

입력

첫째 줄에 정점의 개수 N과 간선의 개수 M, 그리고 철수의 집 번호 K가 주어진다. ( 1 ≤ N, K ≤ 1,000, 1 ≤ M ≤ 100,000 ) 둘째 줄부터 간선의 정보가 주어진다. 각 줄은 두 개의 숫자 a, b, c로 이루어져 있으며, 이는 정점 a에서 정점 b로 이동하는 데 시간이 c만큼 걸린다는 의미이다. 양방향 도로가 아님에 주의하자. 각 정점의 번호는 1번부터 N번까지이다.

출력

마을 사람들이 파티를 즐기기 위해서 이동해야 하는 총 거리의 합을 출력한다.

예제 입력

6 10 3
1 2 1
2 5 2
3 1 5
3 2 3
3 4 1
3 6 4
4 1 1
5 3 1
6 5 3
6 4 2

예제 출력

32

 

 

 

2. 풀이

이 문제는 다익스트라 알고리즘으로 해결할 수 있다. 단, 다익스트라 알고리즘을 통해 예제의 답을 쉽게 구할 수 있어도, 각 노드에서 목표 노드까지의 최단 거리를 구하고 목표 노드에서 각 노드의 최단 거리를 더하는 과정에서 시간 초과 문제가 발생한다.

 

아래의 코드의 경우 일반적으로 다익스트라 알고리즘은 n^2의 시간이 소요되는데, 이를 n번씩 2번 계산을 한다면 어마어마한 계산량이 요구된다.

for(int i = 1; i <= n; i++) {
    sum += dijkstra(i, k);
    sum += dijkstra(k, i);
}

따라서 이 코드를 더 효율적으로 수정해야한다.

 

각 노드에서 목표 노드까지의 거리를 구하는 것은 어쩔 수 없이 위의 코드처럼 반복문을 통해 구해야하지만, 목표 노드에서 각 노드까지의 최단 거리는 반복문을 사용하지 않고 단 한번의 다익스트라 알고리즘을 이용하여 쉽게 구할 수 있다.

 

즉, 목표 노드에서 가장 번호가 높은 노드까지의 최단 거리를 구하면, 목표 노드에서 각 노드까지의 최단 거리를 단 한번의 다익스트라 알고리즘으로 구할 수 있다.

import java.io.*;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.PriorityQueue;
import java.util.StringTokenizer;

class Node implements Comparable<Node>{
    private int node;
    private int weight;

    public Node(int node, int weight) {
        this.node = node; this.weight = weight;
    }

    @Override
    public int compareTo(Node o) {
        return weight - o.weight;
    }

    public int getNode() {
        return this.node;
    }

    public int getWeight() {
        return this.weight;
    }

}

class Graph {

    private ArrayList<ArrayList<Node>> graph;

    public Graph(int nodeSize) {
        graph = new ArrayList<>();

        for(int i = 0; i < nodeSize + 1; i++) {
            graph.add(new ArrayList<>());
        }
    }

    public void put(int n, int m, int weight) {
        graph.get(n).add(new Node(m, weight));
//        graph.get(m).add(new Node(n, weight));
    }

    public ArrayList<Node> getNode(int n) {
        return graph.get(n);
    }
}

public class Main {

    private static final BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    private static final BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));

    private static int n;
    private static int m;
    private static int k;

    private static Graph graph;
    private static int[] dist;
    private static boolean[] isVisited;

    public static void main(String[] args) throws IOException {

        StringTokenizer st = new StringTokenizer(br.readLine());
        n = Integer.parseInt(st.nextToken());
        m = Integer.parseInt(st.nextToken());
        k = Integer.parseInt(st.nextToken());

        graph = new Graph(n);
        dist = new int[n+1];
        isVisited = new boolean[n+1];

        for(int i = 0; i < m; i++) {
            st = new StringTokenizer(br.readLine());
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());
            int c = Integer.parseInt(st.nextToken());

            graph.put(a, b, c);
        }

        int sum = 0;
        for(int i = 1; i <= n; i++) {
            sum += dijkstra(i, k);
        }

        dijkstra(k, n);
        for(int i = 1; i <= n; i++) {
            sum += dist[i];
        }

        bw.write(sum + "");
        br.close();
        bw.flush();
        bw.close();
    }

    private static int dijkstra(int s, int e) {
        Arrays.fill(dist, Integer.MAX_VALUE);
        Arrays.fill(isVisited, false);

        PriorityQueue<Node> pq = new PriorityQueue<>();
        pq.add(new Node(s, 0));
        dist[s] = 0;

        while(!pq.isEmpty()) {
            Node curPoint = pq.poll();
            int curNode = curPoint.getNode();
            int curWeight = curPoint.getWeight();

            if(!isVisited[curNode]) {
                isVisited[curNode] = true;

                for(int i = 0; i < graph.getNode(curNode).size(); i++) {
                    int nextNode = graph.getNode(curNode).get(i).getNode();
                    int nextWeight = graph.getNode(curNode).get(i).getWeight();

                    if(!isVisited[nextNode] && dist[nextNode] > curWeight + nextWeight) {
                        dist[nextNode] = curWeight + nextWeight;
                        pq.add(new Node(nextNode, dist[nextNode]));
                    }
                }
            }
        }

        return dist[e];
    }

}