[백준] 14621번 나만 안되는 연애 C++ 코드

코드

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

using namespace std;

int n, m, lst[1003], u, v, d, cnt, res;
char uv[1003]; // 대학교 성별 저장

vector<pair<int, pair<int, int>>> adj; // 길 저장

// 유니온 파인드 for 크루스칼

int find(int x) {
	if (x == lst[x]) return x;
	return lst[x] = find(lst[x]);
}

void merge(int x, int y) {
	x = find(x);
	y = find(y);
	if (x > y) lst[x] = y;
	else lst[y] = x;
}

// 비용 순 정렬
bool cmp(pair<int, pair<int, int>> a, pair<int, pair<int, int>> b) {
	return a.first < b.first;
}

void FASTIO() {
	ios_base::sync_with_stdio(false);
	cin.tie(NULL);
	cout.tie(NULL);
}

int main() {
	FASTIO();
	cin >> n >> m;
	for (int i = 1; i < n + 1; i++) {
		cin >> uv[i];
		lst[i] = i;
	}
	for (int i = 0; i < m; i++) {
		cin >> u >> v >> d;
		adj.push_back({ d, {u, v} });
	}
	sort(adj.begin(), adj.end(), cmp);
	for (int i = 0; i < adj.size(); i++) {
		int from = adj[i].second.first;
		int to = adj[i].second.second;
		int cost = adj[i].first;
		if (find(from) != find(to) && uv[from] != uv[to]) {
			merge(from, to);
			cnt++;
			res += cost;
		}
	}
    // 길은 n - 1개
    // n - 1가 아닌 경우는 길을 만들 수 없는 경우
	if (cnt != (n - 1)) res = -1;
	cout << res << '\n';
	return 0;
}

풀이

대학교마다 길을 최소비용으로 연결하는 문제이다. 조건이 하나 추가되었는데, 남초는 여초랑 연결, 여초는 남초랑 연결되어야 한다는 점이다. 따라서 크루스칼 알고리즘을 이용하여 두 지점이 연결되었는지 확인할 때 대학교가 다른지도 검사하면 된다.

 

최소스패닝트리에서 길은 n - 1 (n은 노드의 수)개 만들어진다. n - 1개가 아니라면 길을 만들 수 없는 경우이므로 문제에서 준 조건대로 -1을 출력하면 된다.