[백준 / BOJ] C++ 2887 행성 터널

2887번: 행성 터널

 

문제

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

2887번: 행성 터널

 

 

풀이

최소 스패닝 트리(MST)를 구현하기 위해 프림 알고리즘을 사용했다. 각 정점의 3차원 좌표가 주어지고 각 정점 간 실제 비용은 min(|xA-xB|, |yA-yB|, |zA-zB|)이다. 하지만 모든 정점 사이의 |xA-xB|, |yA-yB|, |zA-zB|를 구하게 되면 N^2의 공간복잡도를 가지게 되어 메모리초과가 난다. 따라서 각 좌표에 대해 오름차순으로 정렬한 다음 인접한 정점만 간선으로 추가하게 되면 공간복잡도가 N이 되어 통과할 수 있다. 정렬 시 어떤 정점의 좌표인지 구분하기 위해 pair를 사용해 {좌표, 인덱스}로 벡터에 추가하고 좌표를 기준으로 정렬했다.

 

 

코드

#include <bits/stdc++.h>
using namespace std;
#define pii pair<int, int>

int N;
vector<pii> graph[100001];

// 프림 알고리즘을 사용한 MST
int prim() {
  priority_queue<pii> pq;
  vector<int> visited(N + 1, 0);
  visited[1] = 1;
  for (pii edge : graph[1])
    pq.push(edge);
  int cost = 0, cnt = 1;
  while (cnt < N) {
    auto [weight, node] = pq.top();
    pq.pop();
    if (!visited[node]) {
      visited[node] = 1;
      cost -= weight;
      cnt++;
      for (pii next : graph[node])
        if (!visited[next.second])
          pq.push(next);
    }
  }
  return cost;
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(NULL);
  cin >> N;
  vector<pair<int, int>> x, y, z; // 각 좌표를 저장할 벡터
  for (int i = 1; i <= N; i++) {
    int a, b, c;
    cin >> a >> b >> c;
    // 각 좌표를 정점 번호와 함께 벡터에 저장
    x.push_back({a, i});
    y.push_back({b, i});
    z.push_back({c, i});
  }
  // 좌표를 기준으로 오름차순 정렬
  sort(x.begin(), x.end());
  sort(y.begin(), y.end());
  sort(z.begin(), z.end());
  // 인접한 정점들 간의 거리를 간선으로 추가
  for (int i = 0; i < N - 1; i++) {
    graph[x[i].second].push_back({-abs(x[i].first - x[i + 1].first), x[i + 1].second});
    graph[x[i + 1].second].push_back({-abs(x[i].first - x[i + 1].first), x[i].second});
    graph[y[i].second].push_back({-abs(y[i].first - y[i + 1].first), y[i + 1].second});
    graph[y[i + 1].second].push_back({-abs(y[i].first - y[i + 1].first), y[i].second});
    graph[z[i].second].push_back({-abs(z[i].first - z[i + 1].first), z[i + 1].second});
    graph[z[i + 1].second].push_back({-abs(z[i].first - z[i + 1].first), z[i].second});
  }
  // MST
  cout << prim();
  return 0;
}