[백준] 1368 - 물대기

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

1368번: 물대기

 

처음부터 우물이 파져 있는 것이 아니기 때문에, 약간 까다로운 문제였다.

우물을 파는 방법, 물을 끌어오는 방법 2개가 존재하기 때문에 2가지 방법을 통일하기 위해 가상의 0번째 노드를 추가하였다.

그래서 우물을 파는 방법을 0번째 노드와 연결시켜서, MST를 구현하면 된다.

 

#include <iostream>
#include <algorithm>
#include <queue>
using namespace std;

int n, parent[301];
vector<pair<int, pair<int, int>>> v;

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

void Union(int a, int b) {
	a = Find(a);
	b = Find(b);

	parent[a] = b;
}

int main() {
	cin >> n;

	for (int i = 1; i <= n; i++) {
		int a;
		cin >> a;
		v.push_back({ a,{0,i} });
		parent[i] = i;
	}

	for (int i = 1; i <= n; i++) {
		for (int j = 1; j <= n; j++) {
			int a;
			cin >> a;
			if (i != j) {
				v.push_back({ a, {i,j} });
			}
		}
	}

	int ans = 0;

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

	for (int i = 0; i < v.size(); i++) {
		int cost = v[i].first;
		int p1 = Find(v[i].second.first);
		int p2 = Find(v[i].second.second);

		if (p1 != p2) {
			Union(p1, p2);
			ans += cost;
		}
	}

	cout << ans;

}