[P5] 백준 2887번 행성 터널 C++ MST, 최소 신장 트리

리뷰

 

다른 MST 문제와는 다르게 모든 간선을 구하는 것이 아닌 필요한 간선만 구해 MST를 구하는 문제

기존의 모든 간선을 구하는 방식을 사용했더니 바로 메모리 초과가 출력되었다.

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

 

문제 풀이

전역 변수

• n : 행성의 개수를 저장할 변수
• nodes : Union-Find를 통해 MST를 구할 노드의 배열, 노드의 최대 크기인 10만보다 1크게 설정해 준다.
• Pos : 행성의 각 축 좌표 위치 정보를 나타낼 구조체, sort가 필요하므로 compare함수를 작성해 준다.
• Edge : 행성간 간선의 정보를 나타낼 구조체, 마찬가지로 sort가 필요하다.
• PI : x, y, z축과 행성의 번호를 저장할 Pos타입 벡터

1. n값을 입력 받고, 1 ~ n까지 행성 정보를 입력 받고, nodes배열을 초기화 해준다.
2. 좌표를 입력 받을땐 PI의 0에는 x좌표, 1에는 y좌표, 2에는 z좌표를 행성 번호와 함께 push 해준다.
3. PI의 0 ~ 2인덱스의 벡터를 각 좌표 기준으로 오름차순으로 정렬해 준다. (행성 번호가 뒤죽박죽 되는 것이 정상이다.)
4. 간선 정보를 담을 Edge타입 벡터 edges를 초기화 해준다.
5. x, y, z좌표를 기준으로 가장 인접한 두 행성간 거리의 정보와 행성 번호 정보를 edges에 push해준다.
6. edges를 간선의 가중치를 기준으로 오름차순 정렬해 준다. 이후 Union-Find를 통해 MST를 구해주면 된다.

 

참고 사항

• 기존의 크루스칼을 사용하여 각 행성간 간선을 모두 구해준다면 100000*99999/2 크기가 되어 메모리 초과가 된다.
• 이를 방지하기 위해 행성을 좌표위치순으로 정렬하고 인접한 행성들 끼리의 거리만 구해준 뒤 비교하는 것이다.

 

정답 코드

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

using namespace std;

int n, nodes[100001];

struct Pos {
	int pos, idx;
	bool operator<(const Pos& other) const {
		return pos < other.pos;
	}
};

struct Edge {
	int w, n1, n2;
	bool operator<(const Edge& other) const {
		return w < other.w;
	}
};

vector<Pos> PI[3];

int Find(int a) {
	if (nodes[a] == a) return a;
	return nodes[a] = Find(nodes[a]);
}

bool Union(int a, int b) {
	int A = Find(a);
	int B = Find(b);
	if (A == B) return false;
	nodes[B] = A;
	return true;
}

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

	cin >> n;
	for (int i = 1; i <= n; i++) {
		nodes[i] = i;
		int x, y, z; cin >> x >> y >> z;
		PI[0].push_back({ x, i });
		PI[1].push_back({ y, i });
		PI[2].push_back({ z, i });
	}

	for (int i = 0; i < 3; i++) sort(PI[i].begin(), PI[i].end());

	vector<Edge> edges;
	for (int i = 0; i < n - 1; i++) {
		edges.push_back({ abs(PI[0][i].pos - PI[0][i + 1].pos), PI[0][i].idx, PI[0][i + 1].idx });
		edges.push_back({ abs(PI[1][i].pos - PI[1][i + 1].pos), PI[1][i].idx, PI[1][i + 1].idx });
		edges.push_back({ abs(PI[2][i].pos - PI[2][i + 1].pos), PI[2][i].idx, PI[2][i + 1].idx });
	}

	sort(edges.begin(), edges.end());
	int sum = 0;
	for (Edge edge:edges) {
		int w = edge.w, n1 = edge.n1, n2 = edge.n2;
		if (Union(n1, n2)) sum += w;
	}
	cout << sum;
}