백준 2887 행성 터널

문제

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

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

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

---

코드

import java.io.*;
import java.util.*;
import java.util.stream.IntStream;

public class Main {
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));

        //N 입력 받기
        int n = Integer.parseInt(br.readLine());

        //행성 입력 받기 {인덱스, X, Y, Z} 형태
        int[][] planets=new int[n][4];
        for(int i=0; i<n; i++){
            StringTokenizer st = new StringTokenizer(br.readLine());
            planets[i] = new int[]{i, Integer.parseInt(st.nextToken()), Integer.parseInt(st.nextToken()), Integer.parseInt(st.nextToken())};
        }

        //우선순위 큐 {행성1, 행성2, 거리} 중 거리를 기준으로 최소힙
        PriorityQueue<int[]> pq=new PriorityQueue<>(Comparator.comparingInt(o->o[2]));

        //X, Y, Z 좌표에 대해서 정렬 후 우선순위큐에 삽입
        Arrays.stream(planets)
                .sorted(Comparator.comparingInt(o->o[1]))
                .reduce((o1, o2)-> {
                    pq.add(new int[]{o1[0], o2[0], Math.abs(o1[1]-o2[1])});
                    return o2;
                });
        Arrays.stream(planets)
                .sorted(Comparator.comparingInt(o->o[2]))
                .reduce((o1, o2)-> {
                    pq.add(new int[]{o1[0], o2[0], Math.abs(o1[2]-o2[2])});
                    return o2;
                });
        Arrays.stream(planets)
                .sorted(Comparator.comparingInt(o->o[3]))
                .reduce((o1, o2)-> {
                    pq.add(new int[]{o1[0], o2[0], Math.abs(o1[3]-o2[3])});
                    return o2;
                });

        //부모 배열 초기화
        int[] union=new int[n];
        IntStream.range(0, n).forEach(o->union[o]=o);

        int count=0;
        int answer=0;
        while (!pq.isEmpty()) {
            int[] ptr = pq.poll();

            //부모가 같으면 continue, 다르면 합치기
            if(!combine(ptr[0], ptr[1], union))
                continue;

            //간선 채택
            answer+=ptr[2];
            if(++count==n-1)
                break;
        }

        bw.write(answer+"\n");
        bw.flush();
    }

    //Union Find
    public static boolean combine(int a, int b, int[] union){
        int u1=find(a, union);
        int u2=find(b, union);

        if(u1==u2)
            return false;

        union[u2]=u1;
        return true;
    }

    public static int find(int a, int[] union){
        if(union[a]==a)
            return a;

        return union[a]=find(union[a], union);
    }
}

---

해결 과정

1. 단순 크루스칼 문제인 줄 알았는데, 총 노드 수가 10만 이므로 간선들을 모두 구하면 100억개의 간선이 나온다. 따라서 메모리 초과가 발생했다.

2. 정답을 봤다. 각 좌표별로 정렬을 하면 무조건 시작부터 끝까지의 거리가 해당 좌표의 MST가 된다. 이를 통해서 이웃한 노드끼리의 간선들을 구하면 X좌표의 경우 n-1개가 나오고, Y좌표와 Z좌표 또한 마찬가지이다.
   총 3(n-1)개의 간선을 우선순위큐에 넣게 되므로, 최대 약 30만 개의 간선으로 크루스칼을 진행할 수 있다.

3. Union Find에서 갱신하는 과정을 O(n)으로 구현했더니 시간초과가 발생했다. Memoization을 통한 Union Find를 한다.

4. 간선의 갯수가 많으므로 프림 알고리즘을 통해서도 해보려고 했다. 하지만 임의의 시작점이 끝점이라는 보장이 없으므로 총 길이를 구할 때 구현이 까다로웠다. 프림은 다익스트라처럼 현재까지의 비용을 기준으로 distance 배열을 갱신하는데, 실제 총 길이를 구하려면 채택된 간선들을 더해줘야 하므로 추가적인 구현이 필요했다.

5. 😁