골드3 - 백준 13418 학교 탐방하기

백준 13418 학교 탐방하기

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

접근방법

이 문제는 정렬을 하여 내리막길을 우선해서 MST를 구해주고, 오르막길을 우선해서 MST를 구해주어 각각의 제곱값의 차를 구해주면 되는 문제이다.

풀이

길을 오르막길을 우선해서 먼저 정렬해주었고, 반복문을 순서대로 실행시키면서 compare_union 함수를 통해서 선택된 건물이 연결되어 있는지 확인해주었다. 이미 연결이 되어있다면 건너뛰고 연결이 안 되어있다면 해당 길을 선택하여 cnt변수를 늘리고 building 배열을 갱신해주었다. cnt의 개수로 선택된 오르막길의 수를 구할 수 있다.
다음은 반복문을 거꾸로 실행시켜 같은 방식으로 내리막길을 우선 선택하는 경로를 구해주었고 선택된 오르막길은 result에 개수를 저장해주었다.
마지막으로 각각 선택된 오르막길의 개수를 제곱해주어 차를 구해주었다.

코드

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

using namespace std;

vector<pair<int, pair<int, int>>> roads;
int building[1001];

int compare_union(int num, int change)
{
	if (building[num] == num)
	{
		if (change == -1)
			return num;

		return building[num] = change;
	}

	return building[num] = compare_union(building[num], 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++)
		building[i] = i;

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

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

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

	int cnt = 0;
	for (int i = 0; i < roads.size(); i++)
	{
		int a = roads[i].second.first;
		int b = roads[i].second.second;
		int d = roads[i].first;

		int ra = compare_union(a, -1);
		int rb = compare_union(b, -1);

		if (ra == rb)
			continue;

		compare_union(a, min(ra,rb));
		compare_union(b, min(ra, rb));
		if (d == 0)
			cnt++;
	}

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

	int result = 0;
	for (int i = roads.size()-1; i >= 0; i--)
	{
		int a = roads[i].second.first;
		int b = roads[i].second.second;
		int d = roads[i].first;

		int ra = compare_union(a, -1);
		int rb = compare_union(b, -1);

		if (ra == rb)
			continue;

		compare_union(a, min(ra, rb));
		compare_union(b, min(ra, rb));

		if(d==0)
			result++;
	}

	cout << cnt*cnt - result*result << endl;
	return 0;
}