[백준/BOJ] 백준 17132번 : 두더지가 정보섬에 올라온 이유

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

17132번: 두더지가 정보섬에 올라온 이유

간선을 가중치의 내림차순으로 정렬하여 가중치가 큰 것부터 해당 간선의 가중치가 만족도가 되는 경우의 개수를 찾아서 문제를 풀었다. 이때 확인한 간선의 정점들은 유니온 파인드의 유니온을 하고 해당 정점이 속한 그룹의 개수를 저장하여 이용하는 방식을 통해 문제를 해결했다.

 

코드

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

int n;
vector<pair<int, pair<int, int>>> edge;
long long result = 0;
vector<int> parent(100001);
vector<int> rank_size(100001);
vector<int> group_size(100001);

void Pre()
{
	for (int i = 0; i < 100001; i++)
	{
		parent[i] = i;
		rank_size[i] = 1;
		group_size[i] = 1;
	}
}

int Find(int u)
{
	if (parent[u] == u)
		return u;

	return parent[u] = Find(parent[u]);
}

void Merge(int u, int v)
{
	u = Find(u);
	v - Find(v);

	if (u == v)
		return;

	if (rank_size[u] < rank_size[v])
	{
		parent[u] = v;
		group_size[v] += group_size[u];
	}

	else
	{
		parent[v] = u;
		group_size[u] += group_size[v];

		if (rank_size[u] == rank_size[v])
			rank_size[u]++;
	}
}


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

	Pre();

	cin >> n;

	for (int i = 0; i < n - 1; i++)
	{
		int x, y, w;
		cin >> x >> y >> w;

		edge.push_back(make_pair(w, make_pair(x, y)));
	}

	//가중치의 내림차순으로 정렬
	sort(edge.begin(), edge.end());
	reverse(edge.begin(), edge.end());

	for (int i = 0; i < edge.size(); i++)
	{
		int cost = edge[i].first; //cost는 현재까지 가중치중 가장 작은 가중치가 된다
		int u = edge[i].second.first;
		int v = edge[i].second.second;

		u = Find(u);
		v = Find(v);

		result += ((long long)group_size[u] * (long long)group_size[v] * (long long)cost);
		Merge(u, v);
	}

	cout << result;

	return 0;
}