[BOJ] 2887. 행성 터널

모든 정점을 N-1개로 연결하는 MST 문제이다.

MST를 해결하기 위해선 간선기준인 Kruskal 알고리즘과 정점기준인 Prim 알고리즘을 사용할 수 있다. 두 문제는 어쨋든 출발점, 도착점 비용을 알아야하기 때문에 문제에서 비용을 구하는 방법을 생각해보았다. N이 최대 100,000임으로 2중for문은 사용불가하여 고민해본 결과, 전체 정점을 x,y,z로 각 각 정렬했을때 하나 뒤의 정점과의 x,y,z 차이를 이용하여 간선을 구하는 방법을 고안하였다. 이렇게 모든 간선들의 비용을 PriorityQueue에 저장한다음 UnionFind를 이용하여 MST를 구하였다.

그래프의 간선들을 만듬에 있어 충분한 고민이 필요했던 문제였다.

package boj2;
 
import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.Arrays;
import java.util.Comparator;
import java.util.PriorityQueue;
import java.util.StringTokenizer;
 
public class Main_bj_2887_행성터널 {
    
    static int N;
    static PriorityQueue<Edge> pq;
    static Pos[] arr;
    
    static int[] p;
    
    static int count;
    static int distance;
    
    static class Pos{
        int index,x,y,z;
        Pos(int i,int x,int y,int z){
            index=i;this.x=x;this.y=y;this.z=z;
        }
    }
    
    static class Edge implements Comparable<Edge>{
        int u,v,w;
        Edge(int u,int v,int w) {
            this.u=u;this.v=v;this.w=w;
        }
        
        @Override
        public int compareTo(Edge o) {
            return Integer.compare(w, o.w);
        }
        
    }
    
    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st =null;
        N=Integer.parseInt(br.readLine());
        arr=new Pos[N];
        
        pq=new PriorityQueue<>();
        for(int i=1;i<=N;i++) {
            st=new StringTokenizer(br.readLine());
            int x=Integer.parseInt(st.nextToken());
            int y=Integer.parseInt(st.nextToken());
            int z=Integer.parseInt(st.nextToken());
            arr[i-1]=new Pos(i, x, y, z);
        }
        
        makeEdges();
        
        p=new int[N+1];
        count=0;
        distance=0;
        
        kruskal();
        
        System.out.println(distance);
        
    }
    
    static void makeEdges() {
        Pos start=null;
        Arrays.sort(arr,new Comparator<Pos>() {
            @Override
            public int compare(Pos o1, Pos o2) {
                return Integer.compare(o1.x,o2.x);
            }
        });
        
        start=arr[0];
        for(int i=1;i<N;i++) {
            pq.add(new Edge(start.index, arr[i].index, Math.abs(start.x-arr[i].x)));
            start=arr[i];
        }
        
        Arrays.sort(arr,new Comparator<Pos>() {
            @Override
            public int compare(Pos o1, Pos o2) {
                return Integer.compare(o1.y,o2.y);
            }
        });
        
        start=arr[0];
        for(int i=1;i<N;i++) {
            pq.add(new Edge(start.index, arr[i].index, Math.abs(start.y-arr[i].y)));
            start=arr[i];
        }
        
        Arrays.sort(arr,new Comparator<Pos>() {
            @Override
            public int compare(Pos o1, Pos o2) {
                return Integer.compare(o1.z,o2.z);
            }
        });
        
        start=arr[0];
        for(int i=1;i<N;i++) {
            pq.add(new Edge(start.index, arr[i].index, Math.abs(start.z-arr[i].z)));
            start=arr[i];
        }
    }
    
    static void kruskal() {
        for(int i=1;i<=N;i++)
            p[i]=i;
        
        while(!pq.isEmpty()) {
            Edge edge=pq.poll();
            
            if(find(edge.u)!=find(edge.v)) {    //부모가 다르다면 픽
                union(edge.u, edge.v);
                count+=1;
                distance+=edge.w;
            }
            
            if(count==N-1)
                break;
            
        }
        
    }
    
    static int find(int num) {
        if(p[num]==num)
            return num;
        return p[num]=find(p[num]);
    }
    
    static void union(int num1,int num2) {
        num1=find(num1);
        num2=find(num2);
        
        if(num1<num2)
            p[num2]=num1;
        else
            p[num1]=num2;
    }
}