SWEA 3124 - 최소 스패닝 트리 [Java]

문제

SW Expert Academy

풀이

간선 관리를 쉽게 하기 위해 2차원 배열을 만들기 보다는 Edge 클래스를 선언해서 List로 관리를 했다. 가중치 기준 오름차순 정렬을 위해 Comparable 인터페이스를 implement하고 compareTo 메서드를 오버라이딩했다.

 

처음 제출한 코드가 시간초과에 걸려서 BufferedReader와 BufferedWriter를 쓰기는 했는데, 경로 단축을 안해서 시간 초과가 발생했던 것이라서 그냥 Scanner를 사용해도 되지 않았을까 싶다.

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
 
class Edge implements Comparable<Edge> {
    int weight;
    int from;
    int to;
     
    public Edge(int weight, int from, int to) {
        this.weight = weight;
        this.from = from;
        this.to = to;
    }
 
    @Override
    public int compareTo(Edge o) {
        return this.weight - o.weight;
    }
}
 
public class Solution {
    private static int[] root;
     
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
         
        int T = Integer.parseInt(br.readLine());
        for (int t = 1; t <= T; t++) {
            String[] inputs = br.readLine().split(" ");
            int V = Integer.parseInt(inputs[0]);
            int E = Integer.parseInt(inputs[1]);
             
            List<Edge> edges = new ArrayList<>(E);
            for (int i = 0; i < E; i++) {
                inputs = br.readLine().split(" ");
                int A = Integer.parseInt(inputs[0]);
                int B = Integer.parseInt(inputs[1]);
                int C = Integer.parseInt(inputs[2]);
                edges.add(new Edge(C, A, B));
            }
            Collections.sort(edges);
             
            root = new int[V+1];
            for (int i = 1; i <= V; i++) {
                root[i] = i;
            }
             
            long total = 0;
            int selected = 0;
            for (Edge edge: edges) {
                if (union(edge)) {
                    total += edge.weight;
                    selected++;
                    if (selected == V-1)
                        break;
                }
            }
             
            bw.write("#" + t + " " + total + "\n");
        }
         
        bw.flush();
        bw.close();
        br.close();
    }
     
    private static boolean union(Edge e) {
        int rootFrom = find(e.from);
        int rootTo = find(e.to);
        if (rootFrom == rootTo) return false;
         
        root[rootTo] = rootFrom;
        return true;
    }
     
    private static int find(int x) {
        if (root[x] == x)
            return x;
        else
            return root[x] = find(root[x]);
    }
}