[백준] 2887 행성 터널

2887번: 행성 터널

 

boj.kr

풀이

문제를 읽어보면 MST를 만들어주면 해결할 수 있을거라는 생각이 든다. 그런데 일반적인 문제와 달리 간선의 길이가 주어져있지 않고 min(|xA-xB|, |yA-yB|, |zA-zB|)라는 조건만 나와있다. 이를 위해서 모든 접점 간의 거리를 구해 MST를 만들려고 하면 메모리 초과가 발생한다. 그렇기 때문에 다른 방식을 생각해야 한다.

사실 이 문제는 모든 접점 간의 거리를 구할 필요가 없다. 결국 MST를 만들어야하므로 각 접점 간의 최소 거리가 필요하다. 접점 간의 최소 거리라는 것은 결국 min(|xA-xB|, |yA-yB|, |zA-zB|)을 통해 계산한 거리이다. 이 것을 다시 생각해보면 두 접점 사이에 x값을 기준으로 한 간선, y값을 기준으로 한 간선, z값을 기준으로 한 간선으로 총 3개의 간선이 있다고 볼 수 있다. 따라서 각각의 축을 기준으로 행성들을 정렬한 뒤, 그 중 최소가 되는 간선을 선택해서 MST를 만들어주면 메모리 초과 문제를 해결할 수 있다.

코드

C++

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

struct planet {
    int idx, x, y, z;
};

int n;
vector<planet> p;
vector<pair<int, pair<int, int>>> edges;
vector<int> root;

bool comp_x(planet& p1, planet& p2) {
    return p1.x < p2.x;
}

bool comp_y(planet& p1, planet& p2) {
    return p1.y < p2.y;
}

bool comp_z(planet& p1, planet& p2) {
    return p1.z < p2.z;
}

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

void merge(int a, int b) {
    a = find(a);
    b = find(b);

    if (a < b)
        root[b] = a;
    else
        root[a] = b;
}

int kruskal() {
    int sum = 0;
    for (int i = 0; i < edges.size(); i++) {
        int a = edges[i].second.first;
        int b = edges[i].second.second;
        int c = edges[i].first;
        if (find(a) == find(b)) continue;
        merge(a, b);
        sum += c;
    }
    return sum;
}

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

    cin >> n;
    p = vector<planet>(n);
    root = vector<int>(n);

    for (int i = 0; i < n; i++) root[i] = i;

    for (int i = 0; i < n; i++) {
        cin >> p[i].x >> p[i].y >> p[i].z;
        p[i].idx = i;
    }

    sort(p.begin(), p.end(), comp_x);
    for (int i = 0; i < n - 1; i++)
        edges.push_back({ p[i + 1].x - p[i].x, {p[i].idx, p[i + 1].idx} });

    sort(p.begin(), p.end(), comp_y);
    for (int i = 0; i < n - 1; i++)
        edges.push_back({ p[i + 1].y - p[i].y, {p[i].idx, p[i + 1].idx} });

    sort(p.begin(), p.end(), comp_z);
    for (int i = 0; i < n - 1; i++)
        edges.push_back({ p[i + 1].z - p[i].z, {p[i].idx, p[i + 1].idx} });

    sort(edges.begin(), edges.end());
    cout << kruskal() << '\n';
    return 0;
}

Python3

import sys


class Planet:
    def __init__(self, idx, x, y, z):
        self.idx = idx
        self.x = x
        self.y = y
        self.z = z


def find(a):
    if a != parent[a]:
        parent[a] = find(parent[a])
    return parent[a]


def union(a, b):
    a = find(a)
    b = find(b)

    if a < b:
        parent[b] = a
    else:
        parent[a] = b


def kruskal():
    sum = 0
    for i in edges:
        a = i[1]
        b = i[2]
        c = i[0]
        if find(a) == find(b):
            continue
        union(a, b)
        sum += c
    return sum


n = int(sys.stdin.readline().rstrip())
parent = [0 for _ in range(n)]
p = []
edges = []

for i in range(n):
    parent[i] = i

for i in range(n):
    x, y, z = map(int, sys.stdin.readline().rstrip().split())
    p.append(Planet(i, x, y, z))

p.sort(key=lambda planet: planet.x)
for i in range(n - 1):
    edges.append((p[i + 1].x - p[i].x, p[i].idx, p[i + 1].idx))

p.sort(key=lambda planet: planet.y)
for i in range(n - 1):
    edges.append((p[i + 1].y - p[i].y, p[i].idx, p[i + 1].idx))

p.sort(key=lambda planet: planet.z)
for i in range(n - 1):
    edges.append((p[i + 1].z - p[i].z, p[i].idx, p[i + 1].idx))
edges.sort()
print(kruskal())