[백준] 16202번 MST 게임

문제: 16202번: MST 게임

 

#include <map>
#include <set>
#include <cmath>
#include <queue>
#include <stack>
#include <vector>
#include <string>
#include <iostream>
#include <algorithm>
#include <iomanip> // setprecision을 사용하기 위한 헤더
#include <climits>

using namespace std;

vector<int> nodes;

int find_set(int u) {
	if (nodes[u] == u)
		return u;
	return nodes[u] = find_set(nodes[u]);
}

void union_(int ur, int vr) {
	nodes[vr] = ur;
}

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

	int n, m, k; // 정점의 개수, 그래프 간선의 개수, 턴의 수
	cin >> n >> m >> k;

	nodes = vector<int>(n + 1); // 정점 번호 1부터 시작하길래.. n+1 만큼
	for (int i = 0; i < nodes.size(); i++)
		nodes[i] = i;

	priority_queue< pair<int, pair<int, int>>, vector<pair<int, pair<int, int>>>, greater<pair<int, pair<int, int>>>> minEdges;
	for (int _=0; _<m; _++) {
		// 간선 정보
		int x, y;
		cin >> x >> y;
		minEdges.push(make_pair((_ + 1), make_pair(x, y)));
	}


	int stage = 0;
	bool flag = false; // 한 번 mst 가 안 만들어지기 시작하면 다음은 당연히 안됨. 그거 체크하려고
	while (stage < k) {
		// 복사본 만들기
		priority_queue< pair<int, pair<int, int>>, vector<pair<int, pair<int, int>>>, greater<pair<int, pair<int, int>>>> tmp = minEdges;
		nodes = vector<int>(n + 1); // 정점 번호 1부터 시작하길래.. n+1 만큼
		for (int i = 0; i < nodes.size(); i++)
			nodes[i] = i;


		// stage 만큼 최소 간선 없애기..
		for (int i = 0; i < stage; i++)
			tmp.pop();

		int answer = 0;
		int selectedCnt = 0;
		while ((selectedCnt < n - 1) && (!tmp.empty())) {
			pair<int, pair<int, int>> edge = tmp.top();
			tmp.pop();

			int ur = find_set(edge.second.first);
			int vr = find_set(edge.second.second);

			if (ur != vr) {
				union_(ur, vr);
				selectedCnt++;
				answer += edge.first;
			}
		}

		if (selectedCnt == n - 1)
			cout << answer << " ";
		else 
			flag = true;

		if (flag)
			cout << 0 << " ";

		stage++;
	}


	return 0;
}

 

그냥 기본 MST 알고리즘 코드에다가 각 stage 마다 변경사항 반영해주고 그대로 돌리면 되는 간단한 문제였다.

 

다음 스테이지로 넘어갈 때마다 가중치가 가장 작은 간선을 하나 빼고, 그 상태에서 kruskal 알고리즘을 돌리면 쉽게 풀 수 있다. 

 

아 그리고 매 stage 마다 node 를 새로 초기화 해주는 것도 꼭 기억해야 한다. ^0^