행성 터널

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

문제를 보는 순간 바로 MST문제로 분류할 수 있었다.

하지만, 입력이 사실상 fully connected graph이므로 Prim이건, Kruskal이건 바로는 사용하지 못하는 문제였다.

첨에는 이진 트리를 사용하는 Prim으로 접근했었는데, 좋은 결과를 얻지 못했고,

결국에는 웹 서칭을 해서 아이디어를 얻었다...ㅎㅎ

아리디어의 핵심은 모든 간선을 다 볼 필요 없이 축 별로 인접한 녀석들만 봐주자는 것이다.

아마 구글 검색어 노출도 안되는 내 블로그까지 흘러들어와서 내 글을 보고 있는 당신은 대체 그게 왜 되는데?에 대한 의문을 품고 있으리라.

나 역시 증명에 시간을 꽤 할애하였으므로 증명을 보다 자세히 적어보도록 하겠다.

claim: 모든 간선이 (x축 || y축 || z축)에서 인접한 두 정점으로 이루어진 MST가 존재한다.

proof:

• 귀류법으로 증명한다.
• MST M내에 존재하는 어떠한 간선 (p, q)가 x, y, z 축 어디에서도 인접하지 않았다고 가정하자.
• 일반성을 잃지 않고 w(p, q) == |p.x - q.x|였다고 가정하자.
• 가정에 의해 w(p, r) <= w(p, q)를 만족하는 r이 존재한다.
• M에서 (p, q)를 삭제하면 M은 2개의 트리로 분리된다.
• 역시 일반성을 잃지 않고, r이 q가 속한 트리에 속해있다고 하자.
• 간선 (r, p)를 추가하면 2개의 트리로 분리되었던 M은 다시 하나의 트리가 된다. 이를 M_new이라 하자.
• w(p, r) <= w(p, q)이므로 M_new의 가중치 합이 M보다 작거나 같다.
• 작다면 가정에 모순이고, 같다면 위와 같이 바꾸어 나가면 된다.

#include <iostream>
#include <algorithm>
#include <cstdint>
#include <vector>

using namespace std;

inline int abs(const int x)
{
  return x >= 0 ? x : -x;
}

struct Planet
{
  int id;
  int x, y, z;

  Planet(const int id, const int x, const int y, const int z) : id(id), x(x), y(y), z(z)
  {
  }
};

struct XComparator
{
  bool operator()(const Planet& a, const Planet& b)
  {
    return a.x < b.x;
  }
};

struct YComparator
{
  bool operator()(const Planet& a, const Planet& b)
  {
    return a.y < b.y;
  }
};

struct ZComparator
{
  bool operator()(const Planet& a, const Planet& b)
  {
    return a.z < b.z;
  }
};

struct Edge
{
  int src;
  int dst;
  int weight;
  Edge(const int src, const int dst, const int weight) : src(src), dst(dst), weight(weight)
  {
  }

  bool operator<(const Edge& rhs) const
  {
    return weight < rhs.weight;
  }
};

class DisjointSet
{
private:
  vector<int> parents_;

public:
  DisjointSet(const int size) : parents_(size, -1)
  {
  }

  int Find(const int elem)
  {
    if (parents_[elem] == -1)
    {
      return elem;
    }
    else
    {
      return (parents_[elem] = Find(parents_[elem]));
    }
  }

  void Union(const int a, const int b)
  {
    parents_[Find(a)] = Find(b);
  }
};

int N;
vector<Planet> PLANETS;
vector<Edge> EDGES;

int64_t Kruskal()
{
  int64_t answer = 0;

  DisjointSet dset((int)EDGES.size());
  sort(EDGES.begin(), EDGES.end());

  int num_of_planets = 1;

  for (const auto& edge : EDGES)
  {
    if (num_of_planets == N)
    {
      break;
    }

    if (dset.Find(edge.src) != dset.Find(edge.dst))
    {
      answer += edge.weight;
      dset.Union(edge.src, edge.dst);
      ++num_of_planets;
    }
  }

  return answer;
}

int main(void)
{
  // For faster IO
  ios_base::sync_with_stdio(false);
  cout.tie(nullptr);
  cin.tie(nullptr);

  // Read input
  cin >> N;
  for (int id = 0; id < N; ++id)
  {
    int x, y, z;
    cin >> x >> y >> z;
    PLANETS.emplace_back(id, x, y, z);
  }

  // Extract meaningful edges
  sort(PLANETS.begin(), PLANETS.end(), XComparator());
  for (size_t it = 0; it + 1 < PLANETS.size(); ++it)
  {
    EDGES.emplace_back(PLANETS[it].id, PLANETS[it + 1].id, abs(PLANETS[it].x - PLANETS[it + 1].x));
  }

  sort(PLANETS.begin(), PLANETS.end(), YComparator());
  for (size_t it = 0; it + 1 < PLANETS.size(); ++it)
  {
    EDGES.emplace_back(PLANETS[it].id, PLANETS[it + 1].id, abs(PLANETS[it].y - PLANETS[it + 1].y));
  }

  sort(PLANETS.begin(), PLANETS.end(), ZComparator());
  for (size_t it = 0; it + 1 < PLANETS.size(); ++it)
  {
    EDGES.emplace_back(PLANETS[it].id, PLANETS[it + 1].id, abs(PLANETS[it].z - PLANETS[it + 1].z));
  }

  // Kruskal
  cout << Kruskal() << '\n';
}