백준#2887 행성터널 (java)

전체 코드

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

public class Main {
    public static void main(String[] args) {
        new PlanetTunnel();
    }
}
class PlanetTunnel{
    Buf buf = new Buf();
    int N;
    Planet[] planets;
    PlanetTunnel(){
        PriorityQueue<Planet> xQ = new PriorityQueue<>(Comparator.comparingInt(Planet::getX));
        PriorityQueue<Planet> yQ = new PriorityQueue<>(Comparator.comparingInt(Planet::getY));
        PriorityQueue<Planet> zQ = new PriorityQueue<>(Comparator.comparingInt(Planet::getZ));

        N = buf.getInt();
        planets = new Planet[N];

        for (int i  =0; i<N; i++){
            planets[i] = new Planet(i, buf.getInt(), buf.getInt(), buf.getInt());
            xQ.add(planets[i]);
            yQ.add(planets[i]);
            zQ.add(planets[i]);
        }

        Planet nextXPlanet = xQ.poll();
        Planet nextYPlanet = yQ.poll();
        Planet nextZPlanet = zQ.poll();
        for (int i =0; i<N-1; i++){
            nextXPlanet = linkXPlanet(xQ, nextXPlanet);
            nextYPlanet = linkYPlanet(yQ, nextYPlanet);
            nextZPlanet = linkZPlanet(zQ, nextZPlanet);
        }

        boolean[] isVisited = new boolean[N];
        PriorityQueue<Tunnel> q = new PriorityQueue<>(Comparator.comparingInt(Tunnel::getWeight));
        long result = 0;
        q.addAll(planets[0].getTunnelList());
        isVisited[0] = true;
        int tmp = 0;
        while(tmp!=N-1){
            Tunnel poll = q.poll();
            int idx = poll.linkedPlanet.idx;
            if (!isVisited[idx]){
                result += poll.weight;
                isVisited[idx]=true;
                tmp++;
                q.addAll(planets[idx].tunnelList);
            }
        }
        System.out.println(result);
    }

    public Planet linkXPlanet(PriorityQueue<Planet> q, Planet currentP){
        Planet poll = q.poll();
        int abs = Math.abs(poll.getX() - currentP.getX());
        currentP.addTunnelList(new Tunnel(abs, poll));
        poll.addTunnelList(new Tunnel(abs, currentP));
        return poll;
    }
    public Planet linkYPlanet(PriorityQueue<Planet> q, Planet currentP){
        Planet poll = q.poll();
        int abs = Math.abs(poll.getY() - currentP.getY());
        currentP.addTunnelList(new Tunnel(abs, poll));
        poll.addTunnelList(new Tunnel(abs, currentP));
        return poll;
    }
    public Planet linkZPlanet(PriorityQueue<Planet> q, Planet currentP){
        Planet poll = q.poll();
        int abs = Math.abs(poll.getZ() - currentP.getZ());
        currentP.addTunnelList(new Tunnel(abs, poll));
        poll.addTunnelList(new Tunnel(abs, currentP));
        return poll;
    }
}
class Planet{
    int idx;
    int x;
    int y;
    int z;

    public void addTunnelList(Tunnel tunnel) {
        this.tunnelList.add(tunnel);
    }

    List<Tunnel> tunnelList = new ArrayList<>();

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

    public int getX() {
        return x;
    }

    public int getY() {
        return y;
    }

    public int getZ() {
        return z;
    }

    public List<Tunnel> getTunnelList() {
        return tunnelList;
    }
}
class Tunnel{
    int weight;
    Planet linkedPlanet;

    public Tunnel(int weight, Planet linkedPlanet) {
        this.weight = weight;
        this.linkedPlanet = linkedPlanet;
    }

    public int getWeight() {
        return weight;
    }
}
class Buf{
    BufferedReader br;
    StringTokenizer st;
    Buf() {
        br = new BufferedReader(new InputStreamReader(System.in));
    }
    public String get(){
        while(st==null||!st.hasMoreElements()){
            try{
                st= new StringTokenizer(br.readLine());
            }catch(IOException e){
                return null;
            }
        }
        return st.nextToken();
    }
    public Integer getInt(){
        return Integer.parseInt(get());
    }
}

 

무지성 복붙으로 코드 길이가 길어졌다.

풀이

이 문제를 어떤 가공없이 최소 신장 트리로 푼다면 나올 수 있는 경로가 매우 많아진다. 최대 N이 10만개이므로 (10만*(10만-1))/2개라서  메모리 초과나 시간 초과가 나는 문제이다. 하지만 이 문제의 가중치는 유클리드 거리가 아닌 x, y, z축 하나만 고려하여 가중치를 정할 수 있다. 즉 정렬 각각 정렬해보면 쉽게 풀린다.

 

예시)

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

 

각 행성에 위에부터 A, B..E를 지정했을때 X, Y, Z축에서의 각 위치는 다음과 같다.

x축을 기준으로 A의 위치에서 E로 가는 경로는 차다리 B에서 E로 가는 것이 언제나 빠르며 E로 가기 위해서는 B를 거쳐야만한다. 즉 X, Y, Z는 사실 일방향 트리 구조이다. 따라서 A에게 있어 경로는 저 4가지이다.(곂치니까 사실상 2가지)  따라서 한 행성에서 최대로 가지는 경로는 6가지이므로 이후 MST를 구하면 쉽게 구할 수 있다.

 

구현

PriorityQueue<Planet> xQ = new PriorityQueue<>(Comparator.comparingInt(Planet::getX));
PriorityQueue<Planet> yQ = new PriorityQueue<>(Comparator.comparingInt(Planet::getY));
PriorityQueue<Planet> zQ = new PriorityQueue<>(Comparator.comparingInt(Planet::getZ));
for (int i  =0; i<N; i++){
    planets[i] = new Planet(i, buf.getInt(), buf.getInt(), buf.getInt());
    xQ.add(planets[i]);
    yQ.add(planets[i]);
    zQ.add(planets[i]);
}

Planet nextXPlanet = xQ.poll();
Planet nextYPlanet = yQ.poll();
Planet nextZPlanet = zQ.poll();
for (int i =0; i<N-1; i++){
    nextXPlanet = linkXPlanet(xQ, nextXPlanet);
    nextYPlanet = linkYPlanet(yQ, nextYPlanet);
    nextZPlanet = linkZPlanet(zQ, nextZPlanet);
}

 

 

값을 넣어주면서 정렬하기 위해서 x, y, z축 각각 우선순위 큐를 만들어 삽입해준다. `linkXPlanet`은 다음 `Plant`와 현재 `Plant`의 거리를 계산하고 경로를 생성한다. 즉 가까운 것들끼리 묶는다. 

 

boolean[] isVisited = new boolean[N];
PriorityQueue<Tunnel> q = new PriorityQueue<>(Comparator.comparingInt(Tunnel::getWeight));
long result = 0;
q.addAll(planets[0].getTunnelList());
isVisited[0] = true;
int tmp = 0;
while(tmp!=N-1){
    Tunnel poll = q.poll();
    int idx = poll.linkedPlanet.idx;
    if (!isVisited[idx]){
        result += poll.weight;
        isVisited[idx]=true;
        tmp++;
        q.addAll(planets[idx].tunnelList);
    }
}
System.out.println(result);

 

그 다음 MST 해주면 끝이다.