[C++] 행성 터널 (백준 2887번)

https://www.acmicpc.net/problem/2887

 

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

int n;
vector<vector<int>> a;
vector<tuple<int, int, int>> e;
vector<int> parent;

int find(int a) {
    if (parent[a] == a) return a;
    else return parent[a] = find(parent[a]);
}

bool unite(int a, int b) {
    a = find(a);
    b = find(b);
    if (a != b) {
        parent[a] = b;
        return true;
    } else {
        return false;
    }
}

int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);

    cin >> n;
    
    a = vector<vector<int>>(n, vector<int>(3));
    for (int i = 0; i < n; i++) cin >> a[i][0] >> a[i][1] >> a[i][2];

    for (int i = 0; i < n; i++) {
        for (int j = i + 1;  j < n; j++) {
            int dst = min(min(abs(a[i][0] - a[j][0]), abs(a[i][1] - a[j][1])), abs(a[i][2] - a[j][2]));
            e.emplace_back(dst, i, j);
        }
    }
    
    sort(e.begin(), e.end());

    parent = vector<int>(n);
    for (int i = 0; i < n; i++) parent[i] = i;

    int cnt = 0;
    long long ans = 0;
    for (auto [w, i, j]: e) {
        if (unite(i, j)){
            cnt++;
            ans += w;
        }
        if (cnt == n - 1) break;
    }

    cout << ans << '\n';
}

처음엔 이렇게 가능한 모든 조합에 대해서 간선을 만들어 e 배열에 넣었다. 이렇게 풀면 메모리 초과가 나온다. 모든 쌍의 개수가 10^9 * 10^9 이기 때문이다. (MLE)

 

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

vector<tuple<int, int, int>> e;
vector<int> parent;

int find(int a) {
    if (parent[a] == a) return a;
    else return parent[a] = find(parent[a]);
}

bool unite(int a, int b) {
    a = find(a);
    b = find(b);
    if (a != b) {
        parent[a] = b;
        return true;
    } else {
        return false;
    }
}

int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);

    int n;
    cin >> n;
    
    vector<pair<int, int>> x;
    vector<pair<int, int>> y;
    vector<pair<int, int>> z;

    for (int i = 0; i < n; i++) {
        int tx, ty, tz;
        cin >> tx >> ty >> tz;
        x.emplace_back(tx, i);
        y.emplace_back(ty, i);
        z.emplace_back(tz, i);
    }

    sort(x.begin(), x.end());
    sort(y.begin(), y.end());
    sort(z.begin(), z.end());

    for (int i = 0; i < n - 1; i++) {
        int dst;
        dst = abs(x[i].first - x[i + 1].first);
        e.emplace_back(dst, x[i].second, x[i + 1].second);
        dst = abs(y[i].first - y[i + 1].first);
        e.emplace_back(dst, y[i].second, y[i + 1].second);
        dst = abs(z[i].first - z[i + 1].first);
        e.emplace_back(dst, z[i].second, z[i + 1].second);
    }
    
    sort(e.begin(), e.end());

    parent = vector<int>(n);
    for (int i = 0; i < n; i++) parent[i] = i;

    int cnt = 0;
    long long ans = 0;
    for (auto [w, i, j]: e) {
        if (unite(i, j)){
            cnt++;
            ans += w;
        }
        if (cnt == n - 1) break;
    }

    cout << ans << '\n';
}

잘 생각해보면 가능한 모든 쌍을 전부 볼 필요가 없다. (난 이게 잘 생각이 안 나서 찾아봐서 풀긴 했다.. ㅋㅋ)

한 점 기준으로 가중치가 가장 작은 점의 후보들만 살펴보면 된다. 가중치는 min(x 좌표 차이, y 좌표 차이, z 좌표 차이)이다. 가중치가 x 좌표 차이가 될지, y 좌표 차이가 될지, z 좌표 차이가 될지는 모르지만 셋 중 하나라는 것은 아니까 x 좌표 차이가 가장 작은 점, y 좌표 차이가 가장 작은 점, z 좌표 차이가 가장 작은 점, 이렇게 3개의 점만 확인해주면 된다. 그 중에 무조건 정답이 있을 테니까.

그래서 x 좌표, y 좌표, z 좌표 따로따로 정렬해준 뒤 각각에서 인접한 점들만 가중치를 계산해서 간선 배열 e에 넣어주고, 똑같이 크루스칼 알고리즘으로 MST를 구해주었다. (AC)