백준 14621번, 나만 안되는 연애 풀이

문제

백준 14621번

풀이

빨간색 <-> 파란색을 잇는 간선을 제외하고 모두 제거합니다.

 

남은 간선을 가지고 정점을 모두 연결시킬 수 있는 MST (최소 스패닝 트리)를 만들면 됩니다.

정렬하는데 $ O(MlogM) $, 연결하는데도 $ O(MlogM) $이 걸리니 시간 복잡도 $O(MlogM)$에 해결할 수 있습니다.

 

MST를 구하는 방법 (백준 1197번, 최소 스패닝 트리(MST)를 구하는 크루스칼 알고리즘)

소스코드

#include <bits/stdc++.h>

#define all(v) v.begin(), v.end()

using namespace std;

typedef long long ll;
typedef unsigned long long ull;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef vector<int> vi;
typedef vector<ll> vl;

int UF[1001];
int find(int a) {
	return UF[a] = a == UF[a] ? a : find(UF[a]);
}
int merge(int x, int y) {
	x = find(x), y = find(y);
	if (x != y) {
		UF[x] = y;
		return 1;
	}
	return 0;
}

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

	int n, m; cin >> n >> m;
	vector<char> school(n + 1);
	vector<pair<int, pii>> edge;
	for (int i = 1; i <= n; i++) cin >> school[i];
	while (m--) {
		int u, v, d; cin >> u >> v >> d;
		if (school[u] == school[v]) continue;
		edge.emplace_back(d, pii(u, v));
	}

	for (int i = 0; i <= 1000; i++) UF[i] = i;

	int ans = 0;
	sort(all(edge));
	for (const auto &[d, where] : edge) {
		int x = where.first, y = where.second;
		if (merge(x, y)) {
			ans += d;
		}
	}
	
	int f = 0;
	for (int i = 1; i <= n; i++)
		if (find(i) != find(1)) f = 1;
	
	cout << (!f ? ans : -1);
}