[BOJ] 2887 - 행성 터널

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

BOJ - 행성 터널

문제 유형: MST(최소 신장 트리), 최단 경로, 그래프, 정렬, 이진 탐색

문제 난이도: Platinum V

시간 제한: 1초

메모리 제한: 128MB

 

문제

때는 2040년, 이민혁은 우주에 자신만의 왕국을 만들었다. 왕국은 N개의 행성으로 이루어져 있다. 민혁이는 이 행성을 효율적으로 지배하기 위해서 행성을 연결하는 터널을 만들려고 한다.

행성은 3차원 좌표위의 한 점으로 생각하면 된다. 두 행성 A(xA, yA, zA)와 B(xB, yB, zB)를 터널로 연결할 때 드는 비용은 min(|xA-xB|, |yA-yB|, |zA-zB|)이다.

민혁이는 터널을 총 N-1개 건설해서 모든 행성이 서로 연결되게 하려고 한다. 이때, 모든 행성을 터널로 연결하는데 필요한 최소 비용을 구하는 프로그램을 작성하시오.

 

입력

첫째 줄에 행성의 개수 N이 주어진다. (1 ≤ N ≤ 100,000) 다음 N개 줄에는 각 행성의 x, y, z좌표가 주어진다. 좌표는 -10^9보다 크거나 같고, 10^9보다 작거나 같은 정수이다. 한 위치에 행성이 두 개 이상 있는 경우는 없다. 

 

출력

첫째 줄에 모든 행성을 터널로 연결하는데 필요한 최소 비용을 출력한다.

 

풀이

우선 점들을 입력받으면서, 객체로 저장해놓는다.

그러고, 따로 x좌표들만, y좌표들만, z좌표들만 담는 배열에 담아두고, 대응되는 x,y,z좌표당 점의 id목록들에 현재 점의 id를 넣는다.

 

이제, Prim 알고리즘을 이용해서 풀 수 있다. 0번점을 방문처리한다.

처음 0번 점에서 각각 x,y,z별로 이진탐색을 통해서 각 축기준별로 인덱스를 구한다.

각 축별로 현재 인덱스와 앞 인덱스, 뒤의 인덱스들의 점들은 거리상으로 가까이 있는, 이웃한 점들이라 볼 수 있다.

이 이웃한 점들 중에서 방문처리되지 않은 노드들 각각 (터널 비용, 0, 이웃 노드)의 간선으로 최소 힙 우선순위 큐에 넣는다.

 

이제, 우선순위 큐에서 하나씩 빼보면서, 방문하지 않은 이웃노드라면, 그 이웃노드를 추가하는의미로 방문처리하고, 터널비용을 누적하고, 새로 추가된 노드 기준에서 또 '가까운' 노드들을 최소 힙 우선순위 큐에 추가한다.

이를 n-1번 반복하면 된다. 트리이므로 n-1개의 간선을 형성하기 때문이다.

 

최종적으로 누적된 비용을 출력한다.

 

코드

C++

#include <bits/stdc++.h>

using namespace std;

struct Coordinate {
  int x;
  int y;
  int z;
  Coordinate(int x, int y, int z) {
    this->x = x;
    this->y = y;
    this->z = z;
  }
};

int main(int argc, char *argv[]) {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  cout.tie(0);

  int n;
  cin >> n;
  vector<Coordinate *> arr(n);
  vector<int> xPoses(n);
  vector<int> yPoses(n);
  vector<int> zPoses(n);
  unordered_map<int, vector<int>> xPosToCID;
  unordered_map<int, vector<int>> yPosToCID;
  unordered_map<int, vector<int>> zPosToCID;
  vector<bool> visited(n, false);

  for (int i = 0; i < n; ++i) {
    int x, y, z;
    cin >> x >> y >> z;
    auto C = new Coordinate(x, y, z);
    arr[i] = C;
    xPoses[i] = x;
    yPoses[i] = y;
    zPoses[i] = z;
    xPosToCID[x].emplace_back(i);
    yPosToCID[y].emplace_back(i);
    zPosToCID[z].emplace_back(i);
  }

  sort(xPoses.begin(), xPoses.end());
  sort(yPoses.begin(), yPoses.end());
  sort(zPoses.begin(), zPoses.end());
  xPoses.erase(unique(xPoses.begin(), xPoses.end()), xPoses.end());
  yPoses.erase(unique(yPoses.begin(), yPoses.end()), yPoses.end());
  zPoses.erase(unique(zPoses.begin(), zPoses.end()), zPoses.end());

  priority_queue<tuple<int, int, int>, vector<tuple<int, int, int>>, greater<>>
      pq; // dist, src_id, dst_id
  long long ans = 0;

  auto f = [&](int u) {
    int rankX =
        lower_bound(xPoses.begin(), xPoses.end(), arr[u]->x) - xPoses.begin();
    int rankY =
        lower_bound(yPoses.begin(), yPoses.end(), arr[u]->y) - yPoses.begin();
    int rankZ =
        lower_bound(zPoses.begin(), zPoses.end(), arr[u]->z) - zPoses.begin();

    for (int i = 0; i < xPosToCID[xPoses[rankX]].size(); ++i) {
      int v = xPosToCID[xPoses[rankX]][i];
      if (!visited[v])
        pq.push(make_tuple(0, u, v));
    }

    if (rankX > 0) {
      for (int i = 0; i < xPosToCID[xPoses[rankX - 1]].size(); ++i) {
        int v = xPosToCID[xPoses[rankX - 1]][i];
        int w = min({abs(arr[u]->x - arr[v]->x), abs(arr[u]->y - arr[v]->y),
                     abs(arr[u]->z - arr[v]->z)});
        if (!visited[v])
          pq.push(make_tuple(w, u, v));
      }
    }

    if (rankX < xPoses.size() - 1) {
      for (int i = 0; i < xPosToCID[xPoses[rankX + 1]].size(); ++i) {
        int v = xPosToCID[xPoses[rankX + 1]][i];
        int w = min({abs(arr[u]->x - arr[v]->x), abs(arr[u]->y - arr[v]->y),
                     abs(arr[u]->z - arr[v]->z)});
        if (!visited[v])
          pq.push(make_tuple(w, u, v));
      }
    }

    for (int i = 0; i < yPosToCID[yPoses[rankY]].size(); ++i) {
      int v = yPosToCID[yPoses[rankY]][i];
      if (!visited[v])
        pq.push(make_tuple(0, u, v));
    }

    if (rankY > 0) {
      for (int i = 0; i < yPosToCID[yPoses[rankY - 1]].size(); ++i) {
        int v = yPosToCID[yPoses[rankY - 1]][i];
        int w = min({abs(arr[u]->x - arr[v]->x), abs(arr[u]->y - arr[v]->y),
                     abs(arr[u]->z - arr[v]->z)});
        if (!visited[v])
          pq.push(make_tuple(w, u, v));
      }
    }

    if (rankY < yPoses.size() - 1) {
      for (int i = 0; i < yPosToCID[yPoses[rankY + 1]].size(); ++i) {
        int v = yPosToCID[yPoses[rankY + 1]][i];
        int w = min({abs(arr[u]->x - arr[v]->x), abs(arr[u]->y - arr[v]->y),
                     abs(arr[u]->z - arr[v]->z)});
        if (!visited[v])
          pq.push(make_tuple(w, u, v));
      }
    }

    for (int i = 0; i < zPosToCID[zPoses[rankZ]].size(); ++i) {
      int v = zPosToCID[zPoses[rankZ]][i];
      if (!visited[v])
        pq.push(make_tuple(0, u, v));
    }

    if (rankZ > 0) {
      for (int i = 0; i < zPosToCID[zPoses[rankZ - 1]].size(); ++i) {
        int v = zPosToCID[zPoses[rankZ - 1]][i];
        int w = min({abs(arr[u]->x - arr[v]->x), abs(arr[u]->y - arr[v]->y),
                     abs(arr[u]->z - arr[v]->z)});
        if (!visited[v])
          pq.push(make_tuple(w, u, v));
      }
    }

    if (rankZ < zPoses.size() - 1) {
      for (int i = 0; i < zPosToCID[zPoses[rankZ + 1]].size(); ++i) {
        int v = zPosToCID[zPoses[rankZ + 1]][i];
        int w = min({abs(arr[u]->x - arr[v]->x), abs(arr[u]->y - arr[v]->y),
                     abs(arr[u]->z - arr[v]->z)});
        if (!visited[v])
          pq.push(make_tuple(w, u, v));
      }
    }
  };

  f(0);
  visited[0] = true;

  for (int it = 0; it < n - 1; ++it) {
    while (true) {
      int w, u, v;
      tie(w, u, v) = pq.top();
      pq.pop();

      if (!visited[v]) {
        visited[v] = true;
        ans += w;
        f(v);
        break;
      }
    }
  }

  cout << ans << "\n";

  return 0;
}