LAN

#include <iostream>
#include <vector>
#include <algorithm>
#include <math.h>

#define MAX_V 99999
#define INF 987654321

using namespace std;

double kruskalExtension(vector<pair<int, int>> & selected, vector<pair<int, int>> & preConnected);
int prim(vector<pair<int, int>> & selected);

typedef struct DisjointSet
{
	vector<int> parent, rank;

	DisjointSet(int n) : parent(n), rank(n, 1)
	{
		for (int i = 0; i < n; i++)
			parent[i] = i;
	}

	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[u] > rank[v]) swap(u, v);
		parent[u] = v;

		if (rank[u] == rank[v]) ++rank[v];
	}
} DisjointSet;


int V; // vertex count
vector<pair<int, double>> adj[MAX_V] ;// (v, cost)

int main()
{
	int caseNum;
	cin >> caseNum;

	for (int ci = 0; ci < caseNum; ci++)
	{
		for (int i = 0; i < V; i++)
		{
			adj[i].clear();
		}

		int n, m;
		cin >> n >> m;

		vector<pair<int, int>> pos(n);

		for (int ni = 0; ni < n; ni++)
		{
			int x;
			cin >> x;

			pos[ni].first = x;
		}

		for (int ni = 0; ni < n; ni++)
		{
			int y;
			cin >> y;

			pos[ni].second = y;
		}

		vector<pair<int, int>> connected;
		for (int mi = 0; mi < m; mi++)
		{
			int a, b;
			cin >> a >> b;
			connected.push_back(make_pair(a, b));
		}

		V = n;
		for(int i = 0; i < V; i++)
			for (int j = i + 1; j < V; j++)
			{
				double distance = sqrt(pow(pos[i].first - pos[j].first, 2) + pow(pos[i].second - pos[j].second, 2));

				adj[i].push_back(make_pair(j, distance));
				adj[j].push_back(make_pair(i, distance));
			}
		
		vector<pair<int, int>> selected;

		double minDistance = kruskalExtension(selected, connected);

		cout << "MIN DISTANCE: " << minDistance << endl;
	}
}

double kruskalExtension(vector<pair<int, int>> & selected, vector<pair<int, int>> & preConnected) {
	double ret = 0;
	selected.clear();

	// (cost, (u, v))
	vector<pair<double, pair<int, int>>> edges;

	for (int u = 0; u < V; u++)
	{
		for (int i = 0; i < adj[u].size(); i++)
		{
			int v = adj[u][i].first;

			bool isAlreadyConnected = false;
			for (int pre = 0; pre < preConnected.size(); pre++)
			{
				if (((u == preConnected[pre].first) && (v == preConnected[pre].second)) ||
					((u == preConnected[pre].second) && (v == preConnected[pre].first)))
				{
					isAlreadyConnected = true;
					break;
				}
			}

			if (isAlreadyConnected) continue;

			double cost = adj[u][i].second;

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

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

	DisjointSet sets(V);

	// insert pre-connected vertices
	for (int i = 0; i < preConnected.size(); i++)
	{
		sets.merge(preConnected[i].first, preConnected[i].second);
		selected.push_back(make_pair(preConnected[i].first, preConnected[i].second));

		for (int j = 0; j < adj[preConnected[i].first].size(); j++)
		{
			// find cost of preConnected[i].second
		/*	if (adj[preConnected[i].first][j].first == preConnected[i].second)
			{
				ret += adj[preConnected[i].first][j].second;

				break;
			}*/
		}
	}

	// execute kruskal
	for (int i = 0; i < edges.size(); i++)
	{
		double cost = edges[i].first;
		int u = edges[i].second.first, v = edges[i].second.second;


		// avoid cycle
		if (sets.find(u) == sets.find(v)) continue;

		sets.merge(u, v);

		selected.push_back(make_pair(u, v));
		ret += cost;
	}

	return ret;
}

//int prim(vector<pair<int, int>> & selected) {
//	int ret = 0;
//	selected.clear();
//
//	vector<bool> added(V, false);
//	vector<int> minWeight(V, INF), parent(V, -1);
//
//	minWeight[0] = parent[0] = 0;
//	for (int iter = 0; iter < V; iter++)
//	{
//		int u = -1;
//		for (int v = 0; v < V; v++)
//		{
//			if (!added[v] && (u == -1 || minWeight[u] > minWeight[v]))
//				u = v;
//		}
//
//		if (parent[u] != u)
//			selected.push_back(make_pair(parent[u], u));
//		ret += minWeight[u];
//
//		for (int i = 0; i < adj[u].size(); i++)
//		{
//			int v = adj[u][i].first, weight = adj[u][i].second;
//
//			if (!added[v] && minWeight[v] > weight) {
//				parent[v] = u;
//				minWeight[v] = weight;
//			}
//		}
//	}
//
//	return ret;
//}

 

LanInput.txt

0.00MB

LanOutput.txt

0.00MB

 

이 문제는 기존 Kruskal 알고리즘을 조금 변형하여 이미 몇개의 간선이 선택된 상태에서 Minimum Spanning Tree 를 찾는 문제이다. 알고리즘은 그렇게 어렵지 않아서 금방 구현했는데, 시간을 많이 잡아먹은 부분은 DisjointSet 을 구현하는 부분에서 발생한 버그였다. DisjointSet의 find() 함수에서 return 할때 find(parent[u]) 로 호출해야 하는데 그냥 parent[u] 로 호출해버렸던 것이다. 이로 인해 find() 함수를 호출하면 트리의 parent의 최 상단을 반환하는 것이 아니라 u의 바로 위에 있는 parent만 반환하고 있던 것이었다. 책에 나온 방법은 이미 선택되어 있는 간선의 cost를 0으로 두는 방법인데, 이 방식으로 하면 위의 코드보다 더 간단하게 구현 가능하다.