2020藍橋杯省賽模擬 第九題

問題描述 　　
2015年，全中國實現了戶戶通電。作爲一名電力建設者，小明正在幫助一帶一路上的國家通電。這一次，小明要幫助 n個村莊通電，其中 1 號村莊正好可以建立一個發電站，所發的電足夠所有村莊使用。現在，這 n個村莊之間都沒有電線相連，小明主要要做的是架設電線連接這些村莊，使得所有村莊都直接或間接的與發電站相通。
　　小明測量了所有村莊的位置（座標）和高度，如果要連接兩個村莊，小明需要花費兩個村莊之間的座標距離加上高度差的平方，形式化描述爲座標爲 (x_1, y_1) 高度爲 h_1 的村莊與座標爲 (x_2, y_2) 高度爲 h_2 的村莊之間連接的費用爲sqrt((x_1-x_2)(x_1-x_2)+(y_1-y_2)(y_1-y_2))+(h_1-h_2)*(h_1-h_2)。
　　在上式中 sqrt 表示取括號內的平方根。請注意括號的位置，高度的計算方式與橫縱座標的計算方式不同。
　　由於經費有限，請幫助小明計算他至少要花費多少費用才能使這 n 個村莊都通電。 輸入格式 　　輸入的第一行包含一個整數 n
，表示村莊的數量。 　　接下來 n 行，每個三個整數 x, y, h，分別表示一個村莊的橫、縱座標和高度，其中第一個村莊可以建立發電站。
輸出格式 　　輸出一行，包含一個實數，四捨五入保留 2 位小數，表示答案。
樣例輸入
4 1 1 3 9 9 7 8 8 6 4 5 4
樣例輸出
17.41
評測用例規模與約定 　　
對於 30% 的評測用例，1 <= n <= 10； 　　
對於 60% 的評測用例，1 <= n <= 100； 　　
對於所有評測用例，1 <= n <= 1000，0 <= x, y, h <= 10000。

#include <algorithm>
#include <iostream>
#include <cstring>
#include <string>
#include <vector>
#include <iomanip>
using namespace std;

int n;//邊的數量 
class village;
class edge;
vector<village> v_arr;
vector<edge> e_arr;

class village {
public:
	double x, y, h;
	village(int xx, int yy, int hh)
		:x(xx), y(yy), h(hh) {}
};

class edge {
public:
	int s, e;//邊的連接的兩個村莊的編號-1
	double cost;//連接的花費 
	edge(int start, int end) {
		s = start;
		e = end;
		double x_1 = v_arr[s].x, y_1 = v_arr[s].y, h_1 = v_arr[s].h;
		double x_2 = v_arr[e].x, y_2 = v_arr[e].y, h_2 = v_arr[e].h;
		cost = sqrt((x_1 - x_2) * (x_1 - x_2) + (y_1 - y_2) * (y_1 - y_2)) + (h_1 - h_2) * (h_1 - h_2);
	}
};

double prim(int start) {
	int k = 0;//已連接邊的數量 
	double ans = 0;
	int e_num = e_arr.size();//邊的總數 
	vector<int> e_vis(e_num, 0);//記錄邊是否被訪問 
	vector<int> p_vis(n, 0);//記錄點是否被訪問 
	p_vis[start] = 1;
	while (k < n - 1) {
		double minval = 9999999999999;
		int index=-1;
		//遍歷邊集 
		for (int i = 0; i < e_arr.size(); i++) {
			if (e_vis[i] == 0 && p_vis[e_arr[i].s] + p_vis[e_arr[i].e] == 1) {//0+1||1+0
				if (e_arr[i].cost < minval) {//有更小的花費就更新 
					minval = e_arr[i].cost;
					index = i;
				}
			}
		}
		ans += minval;
		//訪問標記
		e_vis[index] = 1;
		p_vis[e_arr[index].s] = 1;
		p_vis[e_arr[index].e] = 1;
		k++;
	}
	return ans;
}

int main() {
	int x, y, h;
	cin >> n;
	for (int i = 0; i < n; i++) {
		cin >> x >> y >> h;
		v_arr.push_back(village(x, y, h));
	}
	//邊的數量
//	int e_num = (1+(n-1))*(n-1)/2;
	int e_num = 0;
	for (int i = 0; i < n; i++) {
		for (int j = i + 1; j < n; j++) {
			e_arr.push_back(edge(i, j));
		}
	}
	cout << setprecision(2) << fixed << prim(0) << endl;
	return 0;
}