[백준/BOJ] 백준 14621번 : 나만 안되는 연애

https://www.acmicpc.net/problem/14621

14621번: 나만 안되는 연애

같은 성별 학교끼리의 간선은 제외한 그래프로 최소 스패닝 트리를 만든다. 크루스칼 알고리즘을 이용하였다.

 

코드

#include <iostream>
#include <algorithm>
#include <vector>
#include <utility>
using namespace std;

int n, m;
vector<pair<int, int>> adj[1001];
int school[1001];
int parent[1001];
int height[1001];

void Pre()
{
	for (int i = 1; i <= n; i++)
	{
		parent[i] = i;
		height[i] = 0;
	}
}

int find(int node)
{
	if (parent[node] == node)
		return node;

	return parent[node] = find(parent[node]);
}

void merge(int node1, int node2)
{
	node1 = find(node1);
	node2 = find(node2);

	if (node1 == node2)
		return;

	if (height[node1] < height[node2])
		parent[node1] = node2;

	else
	{
		parent[node2] = node1;

		if (height[node1] == height[node2])
			height[node1]++;
	}

}

int Solve()
{
	Pre();

	int ret = 0;
	int cnt = 0;

	vector<pair<int, pair<int, int>>> edge;

	//edge를 구한다
	for (int i = 1; i <= n; i++)
		for (int j = 0; j < adj[i].size(); j++)
		{
			int u = i;
			int v = adj[i][j].second;
			int cost = adj[i][j].first;

			edge.push_back(make_pair(cost, make_pair(u, v)));
		}

	sort(edge.begin(), edge.end());

	for (int i = 0; i < edge.size(); i++)
	{
		int u = edge[i].second.first;
		int v = edge[i].second.second;
		int cost = edge[i].first;

		if (find(u) != find(v))
		{
			merge(u, v);

			ret += cost;
			cnt++;
		}
	}

	//모든 곳을 연결하는 경로가 없을때
	if (cnt != n - 1)
		return -1;

	return ret;
}

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

	char input_c;
	int u, v, d;

	cin >> n >> m;

	for (int i = 1; i <= n; i++)
	{
		cin >> input_c;

		if (input_c == 'M')
			school[i] = 0;
		else
			school[i] = 1;
	}

	for (int i = 0; i < m; i++)
	{
		cin >> u >> v >> d;

		//같은 성별 학교끼리의 간선은 필요 없다
		if (school[u] == school[v])
			continue;

		adj[u].push_back(make_pair(d, v));
		adj[v].push_back(make_pair(d, u));
	}

	cout << Solve();

	return 0;
}