【蓝桥杯/图论】通电 Kruskal解法(Java)

问题描述 　　
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。

import java.util.PriorityQueue;
import java.util.Scanner;
// 1:无需package
// 2: 类名必须Main, 不可修改

public class Main {
    static int n;
    static int[] X,Y,H;
    static int web[];
    static double ans=0;
    static int cnt=0;
    static double dis(int x1,int y1,int x2,int y2,int h1,int h2){
        return Math.sqrt((x1-x2)*(x1-x2)+(y1-y2)*(y1-y2))+(h1-h2)*(h1-h2);
    }
    static class arc implements Comparable<arc>{
        int a,b;
        double dis;
        public arc(int a,int b,double dis){
            this.a=a;
            this.b=b;
            this.dis=dis;
        }
        @Override
        public int compareTo(Main.arc o) {
            return Double.compare(this.dis, o.dis);
        }
    }
    static int find(int x){
        if(web[x]==x)return x;
        return web[x]=find(web[x]);
    }
    static void merge(arc a){
        int fx=find(a.a);
        int fy=find(a.b);
        if(fx!=fy){
            web[fx]=fy;
            ans+=a.dis;
            ++cnt;
        }
        
    }
    static boolean check(){
        return cnt==n-1;
    }
    public static void main(String[] args) {
        Scanner scan = new Scanner(System.in);
        n=scan.nextInt();
        X=new int[n];
        Y=new int[n];
        H=new int[n];
        web=new int[n+1];
        for(int i=0;i<n;i++){
            X[i]=scan.nextInt();
            Y[i]=scan.nextInt();
            H[i]=scan.nextInt();
            web[i+1]=i+1;
        }
        PriorityQueue<arc> queue=new PriorityQueue<arc>();
        for(int i=0;i<n-1;++i){
            for(int j=i+1;j<n;++j){
                arc a=new arc(i,j,dis(X[i],Y[i],X[j],Y[j],H[i],H[j]));
                queue.add(a);
            }
        }
        while(!queue.isEmpty()&&!check()){
            merge(queue.poll());
        }
        System.out.println( String.format("%.2f", ans));
        scan.close();
    }
}