[C++]백준 16562번 : 친구비

문제출처
풀이

'친구의 친구는 친구다'인 상황에서 착안해 같은 set에 속해있는 원소들끼리는 서로 친구라고 생각할 수 있다. 예를 들어 1 3이 친구고 3 5가 서로 친구라면 1 3을 Union하고 3 5를 Union한다. 그럼 {1,3,5}인 set이 만들어지고 1과 5도 친구가 된다. 그리고 이러한 친구 그룹과 친구를 맺을 때 최소비용으로 맺으려면 그룹 내의 친구비 중 최소 친구비를 알아내야한다.

코드

#include<iostream>
#include<vector>
#include<queue>

using namespace std;

int Find(vector<int>& s, int x) {
	if ( s[x] == -1 ) {
		return x;
	}
	return s[x] = Find(s, s[x]);
}
void Union(vector<int> &s, int x, int y , vector<int> &cost) {
	int s_x = Find(s, x);
	int c_x = cost[x];
	cost[s_x] = min(c_x, cost[s_x]);
	int s_y = Find(s, y);
	int c_y = cost[y];
	cost[s_y] = min(c_y, cost[s_y]);
	if ( s_x == s_y ) {
		return;
	}

	if ( s_x < s_y ) {
		s[s_y] = s_x;
		cost[s_x] = min(cost[s_y], cost[s_x]);
	}
	else {
		s[s_x] = s_y;
		cost[s_y] = min(cost[s_y], cost[s_x]);
	}
}

return s[x] = Find(s,s[x]); 을 통해 랭크가 길어지지 않게끔 설계하였다.
Union에서는 cost비교를 하여 가장 적은 cost를 root에 달았다.

int main() {
	ios_base::sync_with_stdio(false); cin.tie(NULL); cout.tie(NULL);
	int N, M, K;
	cin >> N >> M >> K;
	int m = K;
	vector<int> cost(N + 1);
	vector<int> s(N + 1,-1);
	vector<bool> check(N + 1, 0);
	for ( int i = 1; i < N + 1; i++ ) {
		cin >> cost[i];
	}
	for ( int i = 0; i < M; i++ ) {
		int a, b;
		cin >> a >> b;
		Union(s, a, b,cost);
	}

	for ( int i = 1; i < N + 1; i++ ) {
		int root = Find(s,i);
		if ( check[root] == 1 )continue;
		check[root] = 1;
		K -= cost[root];
		if ( K < 0 ) {
			cout << "Oh no" << endl;
			return 0;
		}
	}
	cout << m - K << endl;

	return 0;
}

for문을 통해 그룹내 최소 비용을 구하여 계산하였다.