BOJ 2887 행성 터널

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

문제

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

행성은 3차원 좌표위의 한 점으로 생각하면 된다. 두 행성 A$(x_A, y_A, z_A)$와 B$(x_B, y_B, z_B)$를 터널로 연결할 때 드는 비용은 min$(|x_A-x_B|, |y_A-y_B|, |z_A-z_B|)$이다.

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

---

입력

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

---

출력

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

---

예제 입력 1

5
11 -15 -15
14 -5 -15
-1 -1 -5
10 -4 -1
19 -4 19

예제 출력 1

4

---

코드

처음 문제를 접했을 때는, 다른 최소 스패닝 트리 문제와 별 차이가 없어 보였다.

하지만, 행성들을 잇는 간선을 모두 확인하기엔 너무 수가 많았다. (행성 개수가 최대 100,000개이기 때문에, 확인해야 하는 최대 간선 수는 5,000,050,000)

결국, 다른 방법을 생각해야 했고 모든 행성의 x, y, z 좌표를 각각 오름차순 정렬해서 가장 가까운 행성 사이의 간선만 고려해서 문제를 풀었다.

풀이 과정은 다음과 같다. (예시 케이스는 예제의 x 좌표를 기준으로 설명)

1. x, y, z 좌표를 정렬한다.

   -1 10 11 14 19

2. 간선을 저장하는 edges 우선순위 큐에 가장 가까운 행성의 간선들을 추가한다.

   (-1, 10), (10, 11), (11, 14), (14, 19)의 비용을 구해 간선을 만들어 큐에 추가한다.

3. 간선이 저장된 우선순위 큐에서 크루스칼 알고리즘을 사용하여 최소 스패닝 트리를 만드는 최소 비용을 구한다.

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Comparator;
import java.util.PriorityQueue;
import java.util.StringTokenizer;

public class Main {
    static class Position {
        int planetNumber;
        int position;

        public Position(int planetNumber, int position) {
            this.planetNumber = planetNumber;
            this.position = position;
        }
    }

    static class Edge {
        int start;
        int end;
        int weight;

        public Edge(int start, int end, int weight) {
            super();
            this.start = start;
            this.end = end;
            this.weight = weight;
        }
    }

    static final int DIMENSION = 3;
    static int[] parent;
    static PriorityQueue<Edge> edges = new PriorityQueue<>(Comparator.comparingInt(edge -> edge.weight));

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int N = Integer.parseInt(br.readLine());
        PriorityQueue<Position>[] pos = new PriorityQueue[DIMENSION];
        for (int i = 0; i < DIMENSION; i++) {
            pos[i] = new PriorityQueue<>(Comparator.comparingInt(position -> position.position));
        }

        for (int i = 0; i < N; i++) {
            StringTokenizer st = new StringTokenizer(br.readLine());

            for (int j = 0; j < DIMENSION; j++) {
                int position = Integer.parseInt(st.nextToken());
                pos[j].offer(new Position(i, position));
            }
        }

        Position[] polled = new Position[DIMENSION];
        for (int i = 0; i < DIMENSION; i++) {
            polled[i] = pos[i].poll();
        }

        while (!pos[0].isEmpty()) {
            for (int i = 0; i < DIMENSION; i++) {
                Position peeked = pos[i].peek();
                int distance = Math.abs(polled[i].position - peeked.position);
                Edge newEdge = new Edge(polled[i].planetNumber, peeked.planetNumber, distance);
                edges.offer(newEdge);
                polled[i] = pos[i].poll();
            }
        }

        parent = new int[N];
        for(int i = 0; i < N; i++) {
            parent[i] = i;
        }

        int answer = kruskal();

        System.out.println(answer);
    }

    static int kruskal() {
        int totalCost = 0;

        while (!edges.isEmpty()) {
            Edge edge = edges.poll();

            if(find(edge.start) != find(edge.end)) {
                totalCost += edge.weight;
                union(edge.start, edge.end);
            }
        }

        return totalCost;
    }

    static void union(int parentOfStart, int parentOfEnd) {
        if(parentOfStart > parentOfEnd) {
            parent[parentOfStart] = parentOfEnd;
            return;
        }

        parent[parentOfEnd] = parentOfStart;
    }

    static int find(int planet) {
        if(parent[planet] == planet) {
            return planet;
        }

        return parent[planet] = find(parent[planet]);
    }
}