[백준] 16562번 친구비

문제: 16562번: 친구비

 

#include <map>
#include <set>
#include <cmath>
#include <queue>
#include <stack>
#include <vector>
#include <string>
#include <iostream>
#include <algorithm>

using namespace std;

vector<int> nodes; // 정점들

// 우선 순위 큐가 pair 의 두 번째 요소를 기준으로 최소힙으로 동작할 수 있도록 하는 비교 함수 객체 만들기
struct compare {
	bool operator()(pair<pair<int, int>, int> a, pair<pair<int, int>, int> b) {
		return a.second > b.second;
	}
};

// 각 정점이 속한 집합 번호를 반환하는 함수
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;  // n:학생수, m:친구관계수, k:가진돈
	cin >> n >> m >> k;
	
	// 인덱스 계산 편하게 하려고 n+1 만큼 공간 잡음.
	// 0번 인덱스 요소는 안 쓸 것..
	vector<int> frdValue(n + 1); // 친구비
	for (int i = 1; i < n+1; i++)
		cin >> frdValue[i];

	nodes = vector<int>(n + 1); // 0번은 준석이, 1번부터 n번까지 친구들
	for (int i = 0; i < n + 1; i++)
		nodes[i] = i; // 각 정점이 속한 집합 번호 반영
	

	int selectedEdgeCnt = 0;
	for (int i = 0; i < m; i++) {
		int v, w; // 학생 v, w
		cin >> v >> w;

		int vr = find_set(v);
		int wr = find_set(w);
		
		// 자기자신으로의 연결이랑 중복 연결은 카운트로 안 셀거임
		// 즉, 위 경우가 아닐 때만 연결하고 현재까지 선택된 간선의 수 + 1 해주기
		if (vr != wr) {
			union_(vr, wr);
			selectedEdgeCnt++;
		}

		union_(v, w); // 친구관계 연결해놓기
	}

	// 준석이가 다른 친구들이랑 친구할 때 필요한 비용을 가중치로 생각
	// 우선 순위 큐에는 준석이와 다른 친구랑 직접 친구하는 관계만 넣어놓을 것..
	// 최소힙으로 동작하도록..
	priority_queue<pair<pair<int, int>, int>, vector<pair<pair<int, int>, int>>, compare> edges;
	for (int i = 1; i < n + 1; i++) {
		// '준석이가 i번 친구와 친구하는데 드는 비용' 정보를 담아 간선을 만들기
		pair<pair<int, int>, int> edge = make_pair(make_pair(0, i), frdValue[i]);
		edges.push(edge);
	}

	int sum = 0;
	// 현재 총 노드의 수는 n+1 개임. MST 구할 때 간선은 n개까지만 선택되어야함.
	while (selectedEdgeCnt < n && !edges.empty()) {
		// 일단 가중치가 가장 낮은 간선부터 가져옴
		pair<pair<int, int>, int> minEdge = edges.top();
		edges.pop(); // 가져왔으니까 빼야함

		// 연결되어 있는 경우가 아닐 때만 연결..
		int ur = find_set(minEdge.first.first);
		int vr = find_set(minEdge.first.second);
		if (ur != vr) {
			union_(ur, vr);
			sum += minEdge.second; // 친구비 더해주기
			selectedEdgeCnt++; // 선택된 간선 수 + 1
		}
	}

	// 필요한 친구비 총액이 현재 가진 돈보다 작거나 같으면 가능
	// 돈 부족하면 Oh no 출력..
	if (sum <= k)
		cout << sum << "\n";
	else
		cout << "Oh no\n";


	return 0;
}

 

일단 문제를 어떤식으로 풀었나면 입력으로 받은 친구 관계는 선택된 간선으로 취급하고, 준석이와 나머지 친구들의 직접적인 친구관계를 우선 순위 큐에 넣어놓은 다음 현재 선택된 간선의 개수가 n 보다 작을 때까지 간선을 추가하도록 했다.

 

왜 n 보다 작을 때까지 간선을 추가하도록 해야하냐면 입력으로 n 개의 친구(=정점)가 들어오기는 했지만 준석이까지 정점으로 취급해야 하기 때문에 총 n+1 개의 정점이 있기 때문이다. MST 에서 간선의 개수는 정점의 개수 - 1개이므로 위의 조건을 얻었다.

 

처음에는 입력으로 주어지는 친구 관계를 집합이나 벡터에 넣어 놓은 후 MST 로 처리하는 방법을 생각했는데(깊게 생각하진 않았다. 그래서 쓸 말이 없음) 이보다 더 나은 방법이 떠올랐다. 

 

일단 입력으로 주어지는 친구관계는 무조건 MST 에 포함되어야 한다. 이 조건을 어떤식으로 처리했냐면, 입력으로 친구관계를 받는 시점에 union 을 해준다. 근데 두 정점을 그냥 무조건 합치는 건 안 되고, 자기 자신으로의 연결이랑 중복 연결은 카운트로 안 셀거니까 이 경우를 제외하고 합쳐야 한다. 

 

근데 위와 같은 상황은 그냥 두 정점이 속한 집합의 번호가 다를 경우에 합쳐주는 것과 똑같았다. 그래서 이렇게 합치는 경우에 selectedEdgeCnt 의 값을 +1 해줘서 나중에 MST 계산할 때 반복 횟수를 조정해주었다.

 

이제 다음으로는 준석이랑 다른 친구들의 직접적인 친구관계를 우선 순위 큐에 넣어놨다(최소힙으로 동작하도록). 이렇게 간선들을 다 만든 후 큐에 다 넣어놨으면 본격적으로 Kruskal's Algorithm 을 적용하면 문제가 풀린다. 

 

selectedEdgeCnt 가 n 보다 작을 때까지 반복문을 돌도록 하고, 각 반복문마다 최소 비용 간선을 가져와서 두 정점을 연결하는 것이다. 연결 조건이 만족되면 sum 값의 간선의 가중치(==친구비) 를 더해주고, selectedEdgeCnt 의 값도 +1 해주면 된다. 여기서도 위와 마찬가지로 두 정점이 이미 연결되어 있으면 포기하고 다음 간선으로 넘어간다.

 

맨 마지막으로 해줘야 하는 일은 현재 필요한 총 친구비가 가진 돈보다 작은지 아닌지를 판단해서 출력의 경우를 나눠주는 것이다. 만약 친구비 총액이 현재 가진 돈보다 작거나 같으면 구한 친구비 총액을 출력하도록 하고, 아닌 경우에는 Oh no 를 출력하도록 했다. 

 

근데 나는 이 문제가 MST 로 풀어야 하는 건줄 알았는데 풀고 나서 다른 사람들 풀이 보니까 MST 를 안 쓰고 할 수 있었다. 다른 방법으로 푸는 방식도 알아둬야 할 것 같다.

 

 

 

참고자료:

[C++] 백준/BOJ - 16562 : 친구비

[C++] 백준/BOJ - 16562 : 친구비