골드3 - 백준 14621 나만 안되는 연애

백준 14621 나만 안되는 연애

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

접근방법

남초 학교 여초 학교가 번갈아가면서 이어져야 하기 때문에 성별이 다른 학교만 이어질 수 있도록 MST 조건을 하나 더 추가해주었다.

풀이

node 배열에 집합에 대한 정보를 담는데, 남초 대학인지 여초 대학인지 구분해주었다. 연결 정보를 저장한 route를 정렬한 뒤, 가장 짧은 거리 학교부터 이어주는데, 이때 성별이 같은 학교면 선택하지 않고 넘어가고, 성별이 다른 학교면 union_find 함수를 통해 이미 이어져있는 학교인지 판단해서 차례차례 연결해준다.

코드

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

using namespace std;

int node[1001][2];
bool visited[1001] = {false, };
vector<pair<int, pair<int, int>>> route;

int union_find(int num, int change)
{
	if (node[num][1] == num)
	{
		if (change == -1)
			return num;
		return node[num][1] = change;
	}

	return node[num][1] = union_find(node[num][1], change);
}

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

	int N, M;
	cin >> N >> M;

	for (int i = 0; i <= N; i++)
		node[i][1] = i;

	for (int i = 1; i <= N; i++)
	{
		char s;
		cin >> s;
		if (s == 'M')
			node[i][0] = 0;
		else
			node[i][0] = 1;
	}

	for (int i = 0; i < M; i++)
	{
		int a, b, d;
		cin >> a >> b >> d;

		route.push_back({ d,{ a, b } });
	}

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

	int sum = 0;

	for (int i = 0; i < route.size(); i++)
	{
		int a = route[i].second.first;
		int b = route[i].second.second;
		int dis = route[i].first;

		if (node[a][0] == node[b][0])
			continue;

		int ra = union_find(a, -1);
		int rb = union_find(b, -1);
		if (ra == rb)
			continue;

		sum += dis;
		visited[a] = true;
		visited[b] = true;
		union_find(a, min(ra, rb));
		union_find(b, min(ra, rb));
	}

	for (int i = 1; i <= N; i++)
	{
		if (!visited[i])
		{
			cout << -1 << endl;
			return 0;
		}
	}

	cout << sum << endl;
	return 0;
}