编辑代码

import java.util.*;

// 边结构
class Edge implements Comparable<Edge> {
    int u;
    int v;
    int weight;

    public Edge(int u, int v, int weight) {
        this.u = u;
        this.v = v;
        this.weight = weight;
    }

    @Override
    public int compareTo(Edge other) {
        return Integer.compare(this.weight, other.weight);
    }
}

// 图结构
class Graph {
    private int vertices;
    private List<List<Edge>> adjList;

    public Graph(int vertices) {
        this.vertices = vertices;
        adjList = new ArrayList<>(vertices);
        for (int i = 0; i < vertices; i++) {
            adjList.add(new ArrayList<>());
        }
    }

    public void addEdge(int u, int v, int weight) {
        adjList.get(u).add(new Edge(u, v, weight));
        adjList.get(v).add(new Edge(v, u, weight));
    }

    public void primMST(int start) {
        Edge[] parent = new Edge[vertices];
        int[] key = new int[vertices];
        boolean[] inMST = new boolean[vertices];

        Arrays.fill(key, Integer.MAX_VALUE);
        key[start] = 0;
        parent[start] = null;

        PriorityQueue<Edge> pq = new PriorityQueue<>();
        pq.offer(new Edge(-1, start, 0));

        while (!pq.isEmpty()) {
            Edge current = pq.poll();
            int u = current.v;

            if (inMST[u]) continue;
            inMST[u] = true;

            for (Edge edge : adjList.get(u)) {
                int v = edge.v;
                if (!inMST[v] && edge.weight < key[v]) {
                    key[v] = edge.weight;
                    parent[v] = edge;
                    pq.offer(new Edge(u, v, key[v]));
                }
            }
        }

        printMST(parent);
    }

    private void printMST(Edge[] parent) {
        System.out.println("Minimum Spanning Tree:");
        int totalWeight = 0;
        for (int i = 1; i < vertices; i++) {
            System.out.println(parent[i].u + " - " + parent[i].v + " : " + parent[i].weight);
            totalWeight += parent[i].weight;
        }
        System.out.println("Total Weight: " + totalWeight);
    }
}

// 主程序
public class PrimAlgorithm {
    public static void main(String[] args) {
        Graph graph = new Graph(5);
        graph.addEdge(0, 1, 2);
        graph.addEdge(0, 2, 3);
        graph.addEdge(1, 2, 1);
        graph.addEdge(1, 3, 4);
        graph.addEdge(2, 3, 5);
        graph.addEdge(2, 4, 4);
        graph.addEdge(3, 4, 6);
        graph.primMST(0);
    }
}