BOJ 2887. 행성터널

 

문제 링크 

 

문제 풀이

 이 문제는 $N \leq 10^{5}$개의 행성의 좌표 $x,y,z$가 주어지고 두 행성 $A(x_{A},y_{A},z_{A})$와 $B(x_{B},y_{B},z_{B})$ 의 거리가 $min(|x_{A}-x_{B}|,\space |y_{A}-y_{B}|,\space |z_{A}-z_{B}|)$로 정의될 때 모든 행성들을 연결하는 경로의 최소 길이를 구하는 문제이다.

 

 처음 보자마자 당연하게 MST를 구해야겠다는 생각이 떠오른다. 그런데 MST를 구하는 Kruskal Algorithm을 쓰려고 하면 모든 간선의 개수가 ${N \choose 2}$개이므로 MLE랑 TLE가 날 수밖에 없다. 그래서 간선을 되도록 조금만 저장하는 방향으로 생각해볼 수 있다.

 

 어짜피 MST에 포함되는 간선들은 가중치가 비교적 작은 간선들밖에 없으므로 거리가 짧은 간선들만 고려해주자. 그러기 위해서 일단 두 행성 $A(x_{A},y_{A},z_{A})$와 $B(x_{B},y_{B},z_{B})$ 사이에 간선이 1개라고 생각하지 않고 가중치가 각각 $|x_{A}-x_{B}|,\space |y_{A}-y_{B}|,\space |z_{A}-z_{B}|$ 인 간선 3개가 있다고 생각하자. 이렇게 되면 $x,y,z$를 따로 분리해서 각각끼리만 비교해줄 수 있게 된다.

 

 이제 각 행성마다의 $x,y,z$좌표를 따로 저장을 해주어서 오름차순 정렬을 해주면 각 배열에서 인접한 두 좌표들 사이의 간선이 가중치가 낮다는 것을 알 수 있다. 그래서 $x,y,z$좌표마다 $N-1$개씩, 총 $3N-3$개의 간선들을 이용하여 MST를 구하면 된다. 결론적으로는 $V = N$, $E = 3N-3$이므로 시간복잡도 $O(NlogN)$로 Kruskal Algorithm을 실행할 수 있게 된다.

 

 

코드

#include <iostream>
#include <algorithm>
#include <vector>
using namespace std;
typedef long long ll;
const int NMAX = 100000;

struct point{
    int n;
    ll x,y,z;
};

vector<point> points;
int N,parent[NMAX],rk[NMAX];
vector<pair<int,int>> Xs,Ys,Zs;
vector<pair<ll,pair<int,int>>> edges;
ll MSTweight=0;

void init(){
    for(int i=0;i<N;i++) parent[i] = i;
    sort(Xs.begin(),Xs.end());
    sort(Ys.begin(),Ys.end());
    sort(Zs.begin(),Zs.end());
    for(int i=0;i<N-1;i++){
        pair<int,pair<int,int>> edge1 = {Xs[i+1].first-Xs[i].first,{Xs[i].second,Xs[i+1].second}};
        pair<int,pair<int,int>> edge2 = {Ys[i+1].first-Ys[i].first,{Ys[i].second,Ys[i+1].second}};
        pair<int,pair<int,int>> edge3 = {Zs[i+1].first-Zs[i].first,{Zs[i].second,Zs[i+1].second}};
        edges.push_back(edge1);
        edges.push_back(edge2);
        edges.push_back(edge3);
    }
}

int getparent(int x){
    if(parent[x] == x) return x;
    return parent[x] = getparent(parent[x]);
}

void unionparent(int x, int y){
    x = getparent(x);
    y = getparent(y);
    if(x==y) return;
    if(rk[x] > rk[y]) parent[y] = x;
    else{
        parent[x] = y;
        if(rk[x] == rk[y]) rk[y] ++;
    }
}

void input(){
    cin>>N;
    for(int i=0;i<N;i++){
        int a,b,c;
        cin>>a>>b>>c;
        point p = {i,a,b,c};
        points.push_back(p);
        Xs.push_back({a,i});
        Ys.push_back({b,i});
        Zs.push_back({c,i});
    }
}

void Kruskal(){
    sort(edges.begin(),edges.end());
    for(int i=0;i<edges.size();i++){
        int u = edges[i].second.first;
        int v = edges[i].second.second;
        ll w = edges[i].first;
        if(getparent(u)==getparent(v)) continue;
        unionparent(u,v);
        MSTweight += w;
    }
}

void output(){
    cout<<MSTweight;
}

int main(){
    cin.tie(0);
    cout.tie(0);
    ios::sync_with_stdio(0);
    input();
    init();
    Kruskal();
    output();
}