백준 2887번: 행성 터널 (Java)

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보다 작거나 같은 정수이다. 한 위치에 행성이 두 개 이상 있는 경우는 없다. 

출력

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

 

---

 

풀이

N <= 10만이기에, 가능한 모든 간선은 100억개가 되기에 메모리 초과가 뜬다. 또한, MST가 O(ElogE) 알고리즘이기에 메모리 초과가 (무조건 뜨지만) 안 뜬다고 해도 시간 초과가 뜬다. 따라서, 조금 더 참신한 풀이를 생각할 필요가 있다.

... TreeSet에 넣어서 비교 등의 greedy한 방법만 생각해봤는데 아니고, 생각이 안 나더라. 풀이를 찾아보았다.

어떤 점 u가 있다면, mst를 구성할 때 점 u를 지나는 간선은 u랑 x or y or z좌표가 인접해야 한다. 이러한 간선은 총 6개 있는데, 이를 모두 Edge 클래스의 list에 담는다. 이후 list를 dist에 대해 정렬하여, mst 구성하듯이 하면 된다. 본 풀이에서는 크루스칼을 사용했는데, 프림도 쓸 수 있을 것이다. MST에서 사용될 edge의 수를 미리 줄이는 technique를 배웠다. sort에서 정렬 기준을 부여할 때, 안에 들어가는 변수는 final or effectively final이여야 한다.

 

배운 점

MST에서 사용될 edge의 수를 미리 줄이는 technique를 배웠다. sort에서 정렬 기준을 부여할 때, 안에 들어가는 변수는 final or effectively final이여야 한다.

 

코드

import java.io.*;
import java.util.*;
 
public class Main {
    
    static int[] p;
    static final int[] sort = {1, 2, 3};
    
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        long ans = 0;
        int N = Integer.parseInt(br.readLine());
        List<int[]> points = new ArrayList<>();
        for (int i = 1; i <= N; i++) {
            StringTokenizer st = new StringTokenizer(br.readLine());
            int x = Integer.parseInt(st.nextToken());
            int y = Integer.parseInt(st.nextToken());
            int z = Integer.parseInt(st.nextToken());
            points.add(new int[] {i, x, y ,z});
        }
        
        List<Edge> edges = new ArrayList<>();
        // o1[i]에서 i는 final이여야 한다
        // MST 후보 edge 3*(N-1)개 삽입
        for (int i: sort) {
            Collections.sort(points, (o1, o2) -> {
                return o1[i] - o2[i];
            });
            for (int j = 0; j < points.size()-1; j++) {
                edges.add(new Edge(points.get(j)[0], points.get(j+1)[0], 
                        Math.abs(points.get(j)[i] - points.get(j+1)[i])));
            }
        }
        
        makeSet(N);
        Collections.sort(edges);
        int cnt = 0;
        for (Edge e: edges) {
            if (!union(e.from, e.to)) continue;
            cnt++;
            ans += e.dist;
            if (cnt == N-1) break;
        }
        System.out.println(ans);
    }
    
    static class Edge implements Comparable<Edge> {
        int from;
        int to;
        long dist;
 
        public Edge(int from, int to, long dist) {
            this.from = from;
            this.to = to;
            this.dist = dist;
        }
 
        public int compareTo(Edge o) {
            return Long.compare(this.dist, o.dist);
        }
    }
    
    static void makeSet(int N) {
        p = new int[N + 1];
        for (int i = 0; i <= N; i++) {
            p[i] = i;
        }
    }
 
    static int findRoot(int a) {
        if (a == p[a])
            return a;
        return p[a] = findRoot(p[a]);
    }
 
    static boolean union(int a, int b) {
        a = findRoot(a);
        b = findRoot(b);
        if (a == b)
            return false;
        p[b] = a;
        return true;
    }
}