행성 터널

코드

[ 메모리 초과 코드 ]

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;
int N,ans;
int parent[100002];
struct planet{
    int x;
    int y;
    int z;
};
vector<planet> v;
vector<pair<int, pair<int,int>>> edges;
// findParent()
int findParent(int x){
    if(x == parent[x]) return x;
    return parent[x] = findParent(parent[x]);
}
// unionParent()
void unionParent(int a, int b){
    a = findParent(a);
    b = findParent(b);
    if(a<b) parent[b] = a;
    else parent[a] = b;
}
int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0); cout.tie(0);

    cin >> N;
    for(int i=0;i<N;i++)
    {
        int x, y, z;
        cin >> x >> y >> z;
        planet tmp = {x,y,z}; // 중괄호로 구조체 초기화
        v.push_back(tmp);
    }
    /* 수정해야할 부분 : 모든 경로에 대해 cost를 넣으면 메모리 초과가 발생 */
    for(int i=0;i<v.size();i++)
    {
        for(int j=i;j<v.size();j++)
        {
            int a = i;
            int b = j;
            int cost = min(abs(v[a].x - v[b].x), min(abs(v[a].y - v[b].y), abs(v[a].z - v[b].z)));
            edges.push_back({cost,{a,b}});
        }
    }
    // parent 초기화
    for(int i=1;i<=N;i++) parent[i] = i;
    // edges 정렬
    sort(edges.begin(), edges.end());

    for(int i=0;i<edges.size();i++)
    {
        int a = edges[i].second.first;
        int b = edges[i].second.second;
        int cost = edges[i].first;
        if(findParent(a) != findParent(b)){
            unionParent(a,b);
            ans += cost;
        }
    }
    cout << ans;
    return 0;
}

• 메모리 초과 원인
  • 모든 planet 간의 간선과 비용을 모두 구해서 edges에 넣었기 때문!

[ 정답 코드 ]

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;
long long N,ans;
int parent[100002]; // 1~100000
struct planet{ // 10^9 = 10억
    int idx;
    int x;
    int y;
    int z;
};
vector<planet> v;
vector<pair<long long, pair<int,int>>> edges;
// findParent()
int findParent(int x){
    if(x == parent[x]) return x;
    return parent[x] = findParent(parent[x]);
}
// unionParent()
void unionParent(int a, int b){
    a = findParent(a);
    b = findParent(b);
    if(a<b) parent[b] = a;
    else parent[a] = b;
}
bool compX(planet& n, planet& c){
    return n.x < c.x; // 현재 값이 더 크면 true이니까 바뀜 --> 오름차순!
}
bool compY(planet& n, planet& c){
    return n.y < c.y;
}
bool compZ(planet& n, planet& c){
    return n.z < c.z;
}
int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0); cout.tie(0);

    cin >> N;
    for(int i=0;i<N;i++)
    {
        int x, y, z;
        cin >> x >> y >> z;
        planet tmp = {i+1,x,y,z}; // 중괄호로 구조체 초기화
        v.push_back(tmp);
    }
    /* 모든 경로에 대해 cost를 넣으면 메모리 초과가 발생 
        --> x,y,z좌표축으로 정렬한 후 인접한 행성과의 cost만 edges에 넣기 */
    sort(v.begin(), v.end(), compX);
    for(int i=0;i<v.size()-1;i++) edges.push_back({abs(v[i].x - v[i+1].x),{v[i].idx, v[i+1].idx}});
    sort(v.begin(), v.end(), compY);
    for(int i=0;i<v.size()-1;i++) edges.push_back({abs(v[i].y - v[i+1].y),{v[i].idx, v[i+1].idx}});
    sort(v.begin(), v.end(), compZ);
    for(int i=0;i<v.size()-1;i++) edges.push_back({abs(v[i].z - v[i+1].z),{v[i].idx, v[i+1].idx}});

    // parent 초기화
    for(int i=1;i<=N;i++) parent[i] = i;
    // edges 정렬
    sort(edges.begin(), edges.end());

    for(int i=0;i<edges.size();i++)
    {
        int a = edges[i].second.first;
        int b = edges[i].second.second;
        long long cost = edges[i].first;
        if(findParent(a) != findParent(b)){
            unionParent(a,b);
            ans += cost;
        }
    }
    cout << ans;
    return 0;
}

• 핵심
  • planet간 cost는 min(a.x - b.x, a.y - b.y, a.z - b.z)이다.
  • 즉, x축 / y축 / z축으로 각각 정렬해서 서로 인접한 planet간의 경로만 비교해주면 된다
  • 왜? 어차피 각 축을 기준으로 인접한 행성으로 연결될 수 밖에 없음
    --> MST니까 (최소 신장 트리)