섬 연결하기(문제53)

import java.util.*;

class Solution {
    public int solution(int n, int[][] costs) {
        //MST 포함되는 간선모으기용
        List<Edge> mstEdges = new ArrayList<>();
        //초기화
        PriorityQueue<Edge> pq = new PriorityQueue<>();
        for(int[] cost : costs){
            pq.add(new Edge(cost[0],cost[1],cost[2]));
        }
        DisjointSet disjointSet = new DisjointSet(n);
        while(!pq.isEmpty()){
            Edge edge = pq.poll();
            int rootS = disjointSet.find(edge.start);
            int rootE = disjointSet.find(edge.end);
            
            if(rootS != rootE){
                disjointSet.union(rootS, rootE);
                mstEdges.add(edge);
            }
            
            if(mstEdges.size() == n-1){
                break;
            }
        }
        
        int weightSum = 0;
        for(Edge edge : mstEdges){
            weightSum += edge.weight;
        }
        return weightSum;
    }
    
    class Edge implements Comparable<Edge>{
        int start;
        int end;
        int weight;
        
        public Edge(int s, int e, int w){
            this.start = s;
            this.end = e;
            this.weight = w;
        }
        
        @Override
        public int compareTo(Edge other){
            return this.weight - other.weight;
        }
    }
    
    class DisjointSet{
        int parent[];
        
        public DisjointSet(int n){
            parent = new int[n];
            for(int i=0 ; i<n ; i++){
                parent[i] = i;
            }
        }
        
        public void union(int i, int j){
            int rootI = find(i);
            int rootJ = find(j);
            
            if(rootI != rootJ){
                parent[rootI] = rootJ;
            }
        }
        
        public int find(int i){
            if(parent[i] == i){
                return i;
            }
            return parent[i] = find(parent[i]);        
        }
    }
}

/*
- 모두 이어진거 확인하는 법
- 적은 비용으로 모두 연결
- 음 아무래도 다익스트라......
MST? 
*/