[백준 / 골드4] 1197 최소 스패닝 트리 (Java)

MST를 구하는 알고리즘에는 크루스칼과 프림 알고리즘이 있다.
크루스칼은 O(E log E), 프림은 우선순위 큐를 사용할 시 O(E log V)의 시간 복잡도를 가진다.

문제 보기

사용한 것

• MST를 구하기 위한 프림 알고리즘
• 프림 알고리즘을 구현하기 위한 우선순위 큐
• 간선의 정보를 나타내기 위한 Edge, 우선순위 큐를 사용하기 위해 compareTo() 오버라이딩

풀이 방법

• 모든 간선을 입력받아 저장한 후, 우선순위 큐 pq에 목적지가 1, 비용이 0인 간선을 넣어준다.
• pq가 비거나 모든 정점을 탐색할 때까지 (ct == V) 반복문을 진행한다.
• edge를 pq에서 poll해서 edge의 to를 방문하지 않은 경우에 진행한다. (우선순위 큐에서 뽑아왔으니 가장 비용이 작은 간선)
• 방문 처리 해주고, 비용을 result에 더해준다.
• 이제 해당 정점의 모든 간선들에 대해 방문하지 않은 경우에 pq에 넣어준다. (넣어지면서 가장 비용 적은게 앞으로 감)

코드

public class Main {

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        String[] line = br.readLine().split(" ");
        int V = Integer.parseInt(line[0]);
        List<List<Edge>> graph = new ArrayList<>();
        for (int i = 0; i <= V; i++) {
            graph.add(new ArrayList<>());
        }
        int E = Integer.parseInt(line[1]);
        for (int i = 0; i < E; i++) {
            line = br.readLine().split(" ");
            int v1 = Integer.parseInt(line[0]);
            int v2 = Integer.parseInt(line[1]);
            int cost = Integer.parseInt(line[2]);
            graph.get(v1).add(new Edge(v2, cost));
            graph.get(v2).add(new Edge(v1, cost));
        }

        int ct = 0;
        int result = 0;
        boolean[] visited = new boolean[V + 1];
        PriorityQueue<Edge> pq = new PriorityQueue<>();
        pq.add(new Edge(1, 0));

        while (!pq.isEmpty()) {
            Edge edge = pq.poll();
            if (visited[edge.to]) {
                continue;
            }

            visited[edge.to] = true;
            result += edge.cost;

            for (Edge nextEdge : graph.get(edge.to)) {
                if (!visited[nextEdge.to]) {
                    pq.offer(nextEdge);
                }
            }

            if (++ct == V) {
                break;
            }
        }

        System.out.println(result);
    }
}

class Edge implements Comparable<Edge> {

    int to;
    int cost;

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

    @Override
    public int compareTo(Edge o) {
        return this.cost - o.cost;
    }
}