最小生成树之Kruskal算法

1import java.util.PriorityQueue;
2import java.util.Queue;
3
4class UnionFind {
5
6    private final int[] elements;
7
8    public UnionFind(int size) {
9        elements = new int[size];
10        for (int i = 0; i < elements.length; i++) {
11            elements[i] = i;
12        }
13    }
14
15    public int find(int i) {
16        while (elements[i] != i) {
17            i = elements[i];
18        }
19        return i;
20    }
21
22    public void union(int i, int j) {
23        int a = find(i);
24        int b = find(j);
25        elements[a] = b;
26    }
27}
28
29class Edge implements Comparable<Edge> {
30    public final int i, j, weight;
31
32    public Edge(int i, int j, int weight) {
33        this.i = i;
34        this.j = j;
35        this.weight = weight;
36    }
37
38
39    @Override
40    public int compareTo(Edge o) {
41        return weight - o.weight;
42    }
43}
44
45public class Main {
46
47    private static final int INF = Integer.MAX_VALUE;
48
49    public static void main(String[] args) {
50        int[][] map = {
51                {0, 4, 5, INF, INF, INF},
52                {4, 0, 2, INF, INF, INF},
53                {5, 2, 0, 3, 5, 2},
54                {INF, INF, 3, 0, INF, 1},
55                {INF, INF, 5, INF, 0, 4},
56                {INF, INF, 2, 1, 4, 0}
57        };
58        System.out.println(kruskal(map));
59    }
60
61    public static int kruskal(int[][] map) {
62        /**
63         * 并查集，检测两个点是否已经属于同一棵树
64         */
65        UnionFind unionFind = new UnionFind(map.length);
66        /**
67         * 待选择的边，放入优先队列排序
68         */
69        Queue<Edge> edgeQueue = new PriorityQueue<>();
70        for (int i = 0; i < map.length; i++) {
71            for (int j = 0; j < map[0].length; j++) {
72                if (map[i][j] != INF && i != j) {
73                    edgeQueue.offer(new Edge(i, j, map[i][j]));
74                }
75            }
76        }
77        /**
78         * 总路径长度
79         */
80        int sum = 0;
81        /**
82         * 当前已选择边的数量
83         */
84        int count = 0;
85        while (count < map.length - 1) {
86            Edge edge = edgeQueue.poll();
87            /**
88             * 检测边的两个顶点是否已经在同一棵树中
89             */
90            if (unionFind.find(edge.i) == unionFind.find(edge.j)) {
91                continue;
92            }
93            count++;
94            sum += edge.weight;
95            /**
96             * 将两个顶点所在的树在并查集中合并
97             */
98            unionFind.union(edge.i, edge.j);
99        }
100        return sum;
101    }
102}
103