[백준] 2887 - 행성 터널

https://www.acmicpc.net/problem/2887

2887번: 행성 터널

 

문제

 

때는 2040년, 이민혁은 우주에 자신만의 왕국을 만들었다. 왕국은 N개의 행성으로 이루어져 있다. 민혁이는 이 행성을 효율적으로 지배하기 위해서 행성을 연결하는 터널을 만들려고 한다.

행성은 3차원 좌표위의 한 점으로 생각하면 된다. 두 행성 A(xA, yA, zA)와 B(xB, yB, zB)를 터널로 연결할 때 드는 비용은 min(|xA-xB|, |yA-yB|, |zA-zB|)이다.

 

민혁이는 터널을 총 N-1개 건설해서 모든 행성이 서로 연결되게 하려고 한다. 이때, 모든 행성을 터널로 연결하는데 필요한 최소 비용을 구하는 프로그램을 작성하시오.

 

입력

 

첫째 줄에 행성의 개수 N이 주어진다. (1 ≤ N ≤ 100,000) 다음 N개 줄에는 각 행성의 x, y, z좌표가 주어진다. 좌표는 -109보다 크거나 같고, 109보다 작거나 같은 정수이다. 한 위치에 행성이 두 개 이상 있는 경우는 없다. 

 

출력

 

첫째 줄에 모든 행성을 터널로 연결하는데 필요한 최소 비용을 출력한다.

 

---

풀이

 

MST(Minimum Spanning Tree)를 공부한지 얼마 되지않아 고생 좀 한 문제이다. 검색을 해보니 전부 Kruskal 알고리즘만 사용했길래 Prim으로 도전해보았다!

 

문제의 핵심은 고려할 간선의 개수를 줄이는 것이다. 각 정점 사이의 비용은 기존의 거리 공식이 아닌

행성 A와 B 사이의 비용 = min(|xA-xB|, |yA-yB|, |zA-zB|)

으로 나타내기 때문에 정렬을 통해 각 좌표의 최소값들로 계산해주면 된다.

 

코드로 확인해보면,

// p[n][좌표] : n번째 행성의 x, y, z 좌표와 인덱스를 저장
// x좌표를 기준으로 정렬한 것

Arrays.sort(p, Comparator.comparingInt(o -> o[0]));
for(int i=0; i<n-1; i++) {
    planets.get(p[i][3]).add(new int[]{p[i+1][3], Math.abs(p[i][0] - p[i+1][0])});
    planets.get(p[i+1][3]).add(new int[]{p[i][3], Math.abs(p[i][0] - p[i+1][0])});
}

먼저, 각 좌표별로 정렬을 수행한다. 이후, 저장한 인접 리스트를 사용해서 Prim 알고리즘을 수행한다.

 

// Prim Alogrithm

public static int tunnel(int src) {

    PriorityQueue<int[]> q = new PriorityQueue<>(new Comparator<int[]>() {
        @Override
        public int compare(int[] o1, int[] o2) {
            return o1[1] - o2[1];
	}
    });
    boolean[] visited = new boolean[n];
    int ret = 0;

    for(int[] planet : planets.get(src)) {
        int next = planet[0];
        int nextDistance = planet[1];

    	q.offer(new int[]{next, nextDistance});
    }
    visited[src] = true;

    while(!q.isEmpty()) {
        int now = q.peek()[0];
        int distance = q.peek()[1];
        q.poll();

        if(visited[now]) continue;

        visited[now] = true;
        ret += distance;

        for(int[] planet : planets.get(now)) {
            int next = planet[0];
            int nextDistance = planet[1];

            q.offer(new int[]{next, nextDistance});
    	}
    }
    return ret;
}

PriorityQueue를 이용해서 탐색 시간을 줄일 수 있다.

 

사실, 이 문제의 경우 간선의 수가 많지 않아서 Kruskal 알고리즘이 더 유리할 수 있다. 하지만 사람들이 많이 시도하지 않은 방법으로 해결해보고 싶었다. 아직 Prim 알고리즘도 익숙치 않아서, 간선을 양방향 연결해야하는데 단방향으로 연결하고 맞왜틀을 외쳤다고 한다...

 

 

 

전체코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.*;

public class Main {

    public static int n;
    public static List<List<int[]>> planets;

    public static void input() throws IOException {

        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st;

        n = Integer.parseInt(br.readLine());
        planets = new ArrayList<>();

        int[][] p = new int[n][4];

        for(int i=0; i<n; i++) planets.add(new ArrayList<>());
        for(int i=0; i<n; i++) {
            st = new StringTokenizer(br.readLine(), " ");

            p[i][0] = Integer.parseInt(st.nextToken());
            p[i][1] = Integer.parseInt(st.nextToken());
            p[i][2] = Integer.parseInt(st.nextToken());
            p[i][3] = i;
        }

        Arrays.sort(p, Comparator.comparingInt(o -> o[0]));
        for(int i=0; i<n-1; i++) {
            planets.get(p[i][3]).add(new int[]{p[i+1][3], Math.abs(p[i][0] - p[i+1][0])});
            planets.get(p[i+1][3]).add(new int[]{p[i][3], Math.abs(p[i][0] - p[i+1][0])});
        }

        Arrays.sort(p, Comparator.comparingInt(o -> o[1]));
        for(int i=0; i<n-1; i++) {
            planets.get(p[i][3]).add(new int[]{p[i+1][3], Math.abs(p[i][1] - p[i+1][1])});
            planets.get(p[i+1][3]).add(new int[]{p[i][3], Math.abs(p[i][1] - p[i+1][1])});
        }

        Arrays.sort(p, Comparator.comparingInt(o -> o[2]));
        for(int i=0; i<n-1; i++) {
            planets.get(p[i][3]).add(new int[]{p[i+1][3], Math.abs(p[i][2] - p[i+1][2])});
            planets.get(p[i+1][3]).add(new int[]{p[i][3], Math.abs(p[i][2] - p[i+1][2])});
        }
    }

    public static int tunnel(int src) {

        PriorityQueue<int[]> q = new PriorityQueue<>(new Comparator<int[]>() {
            @Override
            public int compare(int[] o1, int[] o2) {
                return o1[1] - o2[1];
            }
        });
        boolean[] visited = new boolean[n];
        int ret = 0;

        for(int[] planet : planets.get(src)) {
            int next = planet[0];
            int nextDistance = planet[1];

            q.offer(new int[]{next, nextDistance});
        }
        visited[src] = true;

        while(!q.isEmpty()) {
            int now = q.peek()[0];
            int distance = q.peek()[1];
            q.poll();

            if(visited[now]) continue;

            visited[now] = true;
            ret += distance;

            for(int[] planet : planets.get(now)) {
                int next = planet[0];
                int nextDistance = planet[1];

                q.offer(new int[]{next, nextDistance});
            }
        }
        return ret;
    }

    public static void planetTunnel() {

        System.out.println(tunnel(0));
    }

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

        input();
        planetTunnel();
    }
}