Prim算法java实现

原理简单介绍：

1).输入：一个加权连通图，其中顶点集合为V，边集合为E；

2).初始化：Vertex数组，Vertex作为定点，包含其索引，父节点索引，其与已构建的最小生成树中存在的边中最小权值weight，根节点的权重初始为0； 将所有的Vertex加入到一个优先级队列

3).重复下列操作，直到优先级队列为空：

a.在优先级队列里取出最小权值的节点u，将u加入到最小生成树中，并记录u到一个数组res；

b.查找所有u相邻的节点v，如果v不在已生成的最小生成树中，并且u->v的权值比Vertex v的权值小，则更新v的权值

4).输出：res中存储的节点顺序，父节点信息，权值信息就可以完整表达最小生成树。

public class Prim {
    private static class Vertex implements Comparable{
        int index;
        int parent;
        int weight;

        Vertex(int idx) {
            index = idx;
            parent = -1;
            weight = Integer.MAX_VALUE;
        }

        @Override
        public int compareTo(Vertex o) {
            if (weight == o.weight) {
                return 0;
            } else if (weight > o.weight) {
                return 1;
            } else {
                return -1;
            }
        }
    }

    // 返回最小的权重
    public LinkedHashSet prim(int nums, int[][] graph, int root) {
        Vertex[] vertexes = new Vertex[nums + 1]; // 节点编号从1开始，0索引未使用
        HashMap> graphMap = new HashMap<>(); // >
        PriorityQueue priorityQueue = new PriorityQueue<>();
        LinkedHashSet tree = new LinkedHashSet<>(); // 生成的树
        // 初始化顶点，除root外，所有顶点权重为MAX
        for (int i = 1; i <= nums; i++) {
            vertexes[i] = new Vertex(i);
            if (i == root) {
                vertexes[i].weight = 0;
            }
            priorityQueue.add(vertexes[i]);
        }
        // 初始化图
        for (int i = 0; i < graph.length; i++) {
            // u -> v
            HashMap vAndWeight = graphMap.getOrDefault(graph[i][0], new HashMap<>());
            vAndWeight.put(graph[i][1], graph[i][2]);
            graphMap.put(graph[i][0], vAndWeight);
            // v -> u
            vAndWeight = graphMap.getOrDefault(graph[i][1], new HashMap<>());
            vAndWeight.put(graph[i][0], graph[i][2]);
            graphMap.put(graph[i][1], vAndWeight);
        }
        // 开始构建树
        while (!priorityQueue.isEmpty()) {
            Vertex cur = priorityQueue.poll(); // 寻找权重最小的节点
            tree.add(cur);
            // 获取cur的相邻节点
            HashMap map = graphMap.getOrDefault(cur.index, new HashMap<>());
            Iterator> it = map.entrySet().iterator();
            while (it.hasNext()) {
                Map.Entry entry = it.next();
                Vertex vNode = vertexes[entry.getKey()];
                // cur的对端节点没有在tree中，并且cur和对端节点的权值更小
                if (!tree.contains(vNode) && entry.getValue() < vNode.weight) {
                    vNode.parent = cur.index;
                    vNode.weight = entry.getValue();
                    priorityQueue.remove(vNode);
                    priorityQueue.add(vNode);
                }
            }
        }
        return tree;
    }

    public static void main(String[] args) {
        int[][] grap = {
                {1, 2, 2},
                {1, 3, 4},
                {1, 4, 1},
                {2, 4, 3},
                {2, 5, 10},
                {3, 4, 2},
                {3, 6, 5},
                {4, 5, 7},
                {4, 6, 8},
                {4, 7, 4},
                {5, 7, 6},
                {6, 7, 1}
        };
        int n = 7;
        LinkedHashSet actual = new Prim().prim(n, grap, 1);
        Iterator it = actual.iterator();
        while (it.hasNext()) {
            Vertex vertex = it.next();
            System.out.println(vertex.index + ", " + vertex.parent + ", " + vertex.weight);
        }
        System.out.println(actual);
        int[][] grap2 = {
                {1, 2, 2},
                {1, 6, 14},
                {1, 7, 8},
                {2, 3, 19},
                {2, 6, 25},
                {3, 4, 9},
                {3, 5, 5},
                {3, 6, 17},
                {4, 5, 1},
                {5, 6, 13},
                {6, 7, 21}
        };
        actual = new Prim().prim(n, grap2, 1);
        it = actual.iterator();
        while (it.hasNext()) {
            Vertex vertex = it.next();
            System.out.println(vertex.index + ", " + vertex.parent + ", " + vertex.weight);
        }
        System.out.println(actual);
    }
}