【最小生成树】(三) Prim 算法

引入

在最小生成树的前两个章节中, 我们介绍了并查集以及基于并查集的 Kruskal 算法;

不难看出, Kruskal 算法的时间复杂度主要来自于对所有边按权值排序; 假设图中共有 E 条边, 那么 Kruskal 的时间复杂度为 ElogE;

因为我们在并查集中使用了路径压缩算法, 其时间复杂度接近 O(1), 所以整个算法的时间复杂度和排序算法相当, 在使用快速排序的情况下, 就是 ElogE;

如果是稠密图(相对于顶点数量来说, 边的数量多出很多), Kruskal 算法的时间复杂度就会比较高, 这时, 使用 Prim 算法能得到更小的时间复杂度;

Prim 算法也很简单, 选择任意一个顶点开始构建树, 每次选择距离树最近的结点加入树中, 并用新加入的结点, 尝试更新其余未加入结点到树的最小距离; 如此循环, 直到处理完所有顶点;

题目链接: <<寻找宝藏>>

在世界的某个区域，有一些分散的神秘岛屿，每个岛屿上都有一种珍稀的资源或者宝藏。国王打算在这些岛屿上建公路，方便运输。

不同岛屿之间，路途距离不同，国王希望你可以规划建公路的方案，如何可以以最短的总公路距离将 所有岛屿联通起来（注意：这是一个无向图）。

给定一张地图，其中包括了所有的岛屿，以及它们之间的距离。以最小化公路建设长度，确保可以链接到所有岛屿。

解读: 明显, 就是要我们以最小的代价, 将所有点连通, 典型的最小生成树问题;

代码

import java.util.*;

public class Main {
    public static void main(String[] args) {
        Scanner scanner = new Scanner(System.in);
        int v = scanner.nextInt();
        int e = scanner.nextInt();

        // 初始化邻接矩阵，所有值初始化为一个大值，表示无穷大
        int[][] grid = new int[v + 1][v + 1];
        for (int i = 0; i <= v; i++) {
            Arrays.fill(grid[i], 10001);
        }

        // 读取边的信息并填充邻接矩阵
        for (int i = 0; i < e; i++) {
            int x = scanner.nextInt();
            int y = scanner.nextInt();
            int k = scanner.nextInt();
            grid[x][y] = k;
            grid[y][x] = k;
        }

        // 所有节点到最小生成树的最小距离
        int[] minDist = new int[v + 1];
        Arrays.fill(minDist, 10001);

        // 记录节点是否在树里
        boolean[] isInTree = new boolean[v + 1];

        // Prim算法主循环
        for (int i = 1; i < v; i++) {
        	  // 选择距离生成树最近的节点
            int cur = -1;
            int minVal = Integer.MAX_VALUE;
            for (int j = 1; j <= v; j++) {
                if (!isInTree[j] && minDist[j] < minVal) {
                    minVal = minDist[j];
                    cur = j;
                }
            }

            // 将最近的节点加入生成树
            isInTree[cur] = true;

            // 更新非生成树节点到生成树的距离
            for (int j = 1; j <= v; j++) {
                if (!isInTree[j] && grid[cur][j] < minDist[j]) {
                    minDist[j] = grid[cur][j];
                }
            }
        }

        // 统计结果
        int result = 0;
        for (int i = 2; i <= v; i++) {
            result += minDist[i];
        }
        System.out.println(result);
        scanner.close();
    }
}

堆优化代码

如果不采用堆, 时间复杂度为 O(V²), 适用于稠密图;

如果采用堆优化, 时间复杂度就和 Kruskal 算法一致了, 适用于稀疏图, 但代码比 Kruskal 复杂得多, 不如直接用 Kruskal;

private static void prim(){
        Scanner sc = new Scanner(System.in);
        // 顶点数量
        int vNum = sc.nextInt();
        // 边的数量
        int eNum = sc.nextInt();
        // 临接表保存所有边
        List<List<int[]>> graph = new ArrayList<>(vNum + 1);
        for(int i = 0; i <= vNum; i++){
            graph.add(new ArrayList<>());
        }
        for(int i = 0; i < eNum; i++){
            int nodeA = sc.nextInt();
            int nodeB = sc.nextInt();
            int val = sc.nextInt();
            // 因为是无向图, 所以两个方向都要添加;
            // 因为 Prim 算法适用于稠密图, 所以用临接矩阵保存也完全可以, 大家自己可以写邻接矩阵的版本
            graph.get(nodeA).add(new int[]{nodeB, val});
            graph.get(nodeB).add(new int[]{nodeA, val});
        }
        // 记录顶点目前是否在构建的树中; 因为题目规定顶点从1开始, 所以这里多申请一个空间
        boolean[] inTree = new boolean[vNum + 1];
        // 顶点到目前已构建的树的最小距离;
        int[] dist = new int[vNum + 1];
        Arrays.fill(dist, Integer.MAX_VALUE);
        // 采用堆优化, 每次取出当前距离树最近的结点;
        // i[0] = node, i[1] = val;
        PriorityQueue<int[]> heap = new PriorityQueue<>((i1, i2)->Integer.compare(i1[1], i2[1]));
        heap.add(new int[]{1, 0}); // 先把结点1加入最小生成树; 选择任意结点都可以;
        int res = 0; // 记录结果
        // 只需要选择 n - 1条边;
        for(int i = 0; i < vNum -1; i++){
            int[] curNode = heap.poll();
            // 当我们更新一个顶点到树的最小距离时, 直接将更新后的结果加入了堆中, 没有删除原来的;
            // 所以可能出现顶点已经处理过了, 后面又再次遇到的情况; 需要跳过;
            if(inTree[curNode[0]]){
                continue;
            }
            inTree[curNode[0]] = true;
            res += curNode[1];
            // 对于还未加入树的结点, 考虑由curNode 连接到树的代价, 是否更小了, 如果是, 更新最小距离;
            List<int[]> edges = graph.get(curNode[0]);
            for(int[] edge : edges){
                if(!inTree[edge[0]] && edge[1] < dist[edge[0]]){
                    dist[edge[0]] = edge[1];
                    heap.add(new int[]{edge[0], edge[1]});
                }
            }
        }
    }