[백준 10423] 전기가 부족해 (C++)

#include <bits/stdc++.h>
#define FASTIO ios::sync_with_stdio(false); cin.tie(0); cout.tie(0);
#define INF 987654321
using namespace std;
//루트 노드를 찾아주는 코드
typedef pair <int, float>pii;


vector<int> city;
vector <int> gen; //발전기 있는 도시 
int parent[5052];
int N,M,K; 

struct st {
	int s, e;
	float c;
};
//간선간의 cost를 저장해 주자.
st arr[100001];

bool compare(st a, st b) {
	return a.c < b.c;
}

int find(int a) {
	//cout << "find a =" << a<<"\n";
	//a가 root임
	if (parent[a] < 0) return a;
	//a가 root가 아닐때 
	return parent[a] = find(parent[a]);
}

int merge(int a, int b) {
	//cout << "merge a=" << a << " b =" << b << "\n";
	bool connect_a =false, connect_b = false;
	a = find(a);
	b = find(b);
	//cout << "merge a.patent=" << a << " b.parent =" << b << "\n";
	if (a == b)return 0;
	//발전소가 parnet 이면 pass 하자 .ㅇㄸ ? 
	for (int i = 0; i < K ; i++) {
		//cout << "gem[" << i << "] =" << gen[i] << "\n";
		if (a == gen[i])
			connect_a = true;
		if (b == gen[i])
			connect_b = true;	
	}
	if (connect_a == true && connect_b == true)
		return 0;

	//small to large
	if (connect_a == true) {
		parent[a] += parent[b];
		parent[b] = a;
		return 1;
	}
	if (connect_b == true) {
		parent[b] += parent[a];
		parent[a] = b;
		return 1;
	}
	if (parent[a] > parent[b]) swap(a, b);
	parent[a] += parent[b];
	parent[b] = a;
	return 1;
}

void slove() {
	//간선이 작은 순으로 정렬
	sort(arr, arr + M, compare);
	float ans = 0;
	for (int i = 1; i <= N; i++) parent[i] = -1;

	for (int i = 0; i < M; i++) {
		if (merge(arr[i].s, arr[i].e)) {
			ans += arr[i].c;
		}
	}
	cout << ans ;
}


int main() {
	FASTIO;
	int a, b, c;
	cin >> N >> M >>K;
	//cout << N << M << K<<"\n";
	for (int i = 0; i < K; i++) {
		//K개의 발전기 도시 입력 받기 
		cin >> a; 
		gen.push_back(a);
	}
	for (int i = 0; i < M; i++) {
		// 도시간의 cost 입력받기
		cin >> a >> b >> c;
		arr[i].s = a; arr[i].e = b; arr[i].c = c;
	}

	slove();
}

union find를 이용한 크루스칼 알고리즘을 이용해 mst 를 만드는것을 응용해서 풀리는 문제이다.

 

응용해야 하는 부분은 , merge 함수 부분이다. 

merge 해야하는 두 노드의 parent가 발전소라면, 두 노드는 더이상 merge 하지 않는다.

 

그리고 merge 할 때 하나의 노드가발전기라면 , 그 노드에 merge 해 준다.

 

두 노드 모두 발전기가 아니라면 두 그냥 더 큰 노드를 parent로 설정해 준다. 

int merge(int a, int b) {
	//cout << "merge a=" << a << " b =" << b << "\n";
	bool connect_a =false, connect_b = false;
	a = find(a);
	b = find(b);
	//cout << "merge a.patent=" << a << " b.parent =" << b << "\n";
	if (a == b)return 0;
	//발전소가 parnet 이면 pass 하자 . 
	for (int i = 0; i < K ; i++) {
		//cout << "gem[" << i << "] =" << gen[i] << "\n";
		if (a == gen[i])
			connect_a = true;
		if (b == gen[i])
			connect_b = true;	
	}
	if (connect_a == true && connect_b == true)
		return 0;

	//small to large
	if (connect_a == true) {
		parent[a] += parent[b];
		parent[b] = a;
		return 1;
	}
	if (connect_b == true) {
		parent[b] += parent[a];
		parent[a] = b;
		return 1;
	}
	if (parent[a] > parent[b]) swap(a, b);
	parent[a] += parent[b];
	parent[b] = a;
	return 1;
}