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

 

출력

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

 

LOGIC

1. 터널로 연결할 때 드는 비용은 min(|xA-xB|, |yA-yB|, |zA-zB|) 이므로 x기준, y기준, z기준으로 모두 정렬하여 Heap에 저장한다.
2. Heap에서 차례데로 행성들을 연결하며 간선의 길이가 N-1이 되면 MST가 완성된다.
3. 정렬된 가중치를 기준으로 MST를 만드므로 KRUSKAL 알고리즘을 사용한다.

행성을 객체로 표현하기 위한 Class를 선언하는데 입력은 x, y, z좌표만 주어지므로 From, To를 구현하기 위해 입력 순서대로 행성에 번호를 매긴다. 번호를 표현하기 위한 idx 변수를 추가한다.

static class Planet {
    int x, y, z, idx;

    public Planet(int x, int y, int z, int idx) {
        this.x = x;
        this.y = y;
        this.z = z;
        this.idx = idx;
    }
}

 

행성간의 연결을 나타내는 터널 객체를 위한 Class를 선언하는데, 가중치를 기준으로 정렬해야 하므로 Comparable를 구현한다.

static class Tunnel implements Comparable<Tunnel> {
    int from, to, weight;

    public Tunnel(int from, int to, int weight) {
        this.from = from;
        this.to = to;
        this.weight = weight;
    }

    @Override
    public int compareTo(Tunnel o) {
        return Integer.compare(this.weight, o.weight);
    }
}

 

행성을 입력받으며 서로소 집합을 표현하기 위한 Parents 배열을 같이 초기화 후 입력 받는다.

planet = new Planet[N];
parents = new int[N];
for (int i = 0; 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());

    planet[i] = new Planet(x, y, z, i);
    parents[i] = i;
}

 

행성을 모두 입력 받으면 터널을 생성한다. Logic대로 x, y, z 기준으로 모두 Heap에 담는다. 그러면 행성간 최소 거리 순서대로 Heap에 쌓이게 된다.

PriorityQueue<Tunnel> pq = new PriorityQueue<>();

//x기준 정렬 후 pq에 삽입
Arrays.sort(planet, new Comparator<Planet>() {
	@Override
	public int compare(Planet o1, Planet o2) {
		return o1.x - o2.x;
	}
});
for (int i = 0; i < N-1; i++) {
	pq.add(new Tunnel(planet[i].idx, planet[i+1].idx, planet[i+1].x - planet[i].x));
}

//y기준 정렬 후 pq에 삽입
Arrays.sort(planet, new Comparator<Planet>() {
	@Override
	public int compare(Planet o1, Planet o2) {
		return o1.y - o2.y;
	}
});
for (int i = 0; i < N-1; i++) {
	pq.add(new Tunnel(planet[i].idx, planet[i+1].idx, planet[i+1].y - planet[i].y));
}

//z기준 정렬 후 pq에 삽입
Arrays.sort(planet, new Comparator<Planet>() {
	@Override
	public int compare(Planet o1, Planet o2) {
		return o1.z - o2.z;
	}
});
for (int i = 0; i < N-1; i++) {
	pq.add(new Tunnel(planet[i].idx, planet[i+1].idx, planet[i+1].z - planet[i].z));
}

 

이후 Heap에서 순서대로 꺼내며 연결되지 않은 행성일 경우 간선의 수가 N-1일 때까지 연결을 시키며 연결된 거리를 누적시킨다. 연결은 KRUSKAL 알고리즘대로 UNION-FIND를 사용한다.

int result = 0, cnt = 0;
while(!pq.isEmpty()) {
	Tunnel tn = pq.poll();
	
	if(union(tn.to, tn.from)) {
		cnt++;
		result += tn.weight;
	}
	
	if(cnt == N-1) break;
}

---

CODE

package P2887;

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

/**
 * 백준 2887. 행성 터널- 플래티넘 5
 * @category MST
 */
public class Main {
	static class Planet {
		int x, y, z, idx;

		public Planet(int x, int y, int z, int idx) {
			this.x = x;
			this.y = y;
			this.z = z;
			this.idx = idx;
		}
	}
	
	static class Tunnel implements Comparable<Tunnel> {
		int from, to, weight;
	
		public Tunnel(int from, int to, int weight) {
			this.from = from;
			this.to = to;
			this.weight = weight;
		}

		@Override
		public int compareTo(Tunnel o) {
			return Integer.compare(this.weight, o.weight);
		}
	}
	
	static int N;
	static Planet[] planet;
	static int[] parents;

	public static void main(String[] args) throws Exception {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer st;

		N = Integer.parseInt(br.readLine());
		planet = new Planet[N];
		parents = new int[N];
		for (int i = 0; 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());
			
			planet[i] = new Planet(x, y, z, i);
			parents[i] = i;
		}
		
		PriorityQueue<Tunnel> pq = new PriorityQueue<>();
		
		//x기준 정렬 후 pq에 삽입
		Arrays.sort(planet, new Comparator<Planet>() {
			@Override
			public int compare(Planet o1, Planet o2) {
				return o1.x - o2.x;
			}
		});
		for (int i = 0; i < N-1; i++) {
			pq.add(new Tunnel(planet[i].idx, planet[i+1].idx, planet[i+1].x - planet[i].x));
		}
		
		//y기준 정렬 후 pq에 삽입
		Arrays.sort(planet, new Comparator<Planet>() {
			@Override
			public int compare(Planet o1, Planet o2) {
				return o1.y - o2.y;
			}
		});
		for (int i = 0; i < N-1; i++) {
			pq.add(new Tunnel(planet[i].idx, planet[i+1].idx, planet[i+1].y - planet[i].y));
		}
		
		//z기준 정렬 후 pq에 삽입
		Arrays.sort(planet, new Comparator<Planet>() {
			@Override
			public int compare(Planet o1, Planet o2) {
				return o1.z - o2.z;
			}
		});
		for (int i = 0; i < N-1; i++) {
			pq.add(new Tunnel(planet[i].idx, planet[i+1].idx, planet[i+1].z - planet[i].z));
		}
		
		//연결된 간선의 수가 N-1개일 때 까지 pq에서 꺼내며 union한다.
		int result = 0, cnt = 0;
		while(!pq.isEmpty()) {
			Tunnel tn = pq.poll();
			
			if(union(tn.to, tn.from)) {
				cnt++;
				result += tn.weight;
			}
			
			if(cnt == N-1) break;
		}
		
		System.out.println(result);
	}
	
	static boolean union(int a, int b) {
		int ar = find(a);
		int br = find(b);
		
		if(ar == br) return false;
		
		parents[br] = ar;
		return true;
	}
	
	static int find(int a) {
		if(parents[a] == a) return a;
		
		return parents[a] = find(parents[a]);
	}
}