用 java语言以测试驱动开发实现最小生成树算法：Prim 和 Kruskal 算法

Prim算法：

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.PriorityQueue;
import java.util.Set;

public class PrimMST {

    public static List<Edge> primMST(Graph graph) {
        List<Edge> result = new ArrayList<>();
        Set<Integer> visited = new HashSet<>();
        PriorityQueue<Edge> pq = new PriorityQueue<>((a, b) -> a.weight - b.weight);

        visited.add(0);
        for (Edge e : graph.edges[0]) {
            pq.offer(e);
        }

        while (!pq.isEmpty() && visited.size() < graph.vertices) {
            Edge e = pq.poll();
            if (visited.contains(e.to)) {
                continue;
            }

            visited.add(e.to);
            result.add(e);

            for (Edge next : graph.edges[e.to]) {
                if (!visited.contains(next.to)) {
                    pq.offer(next);
                }
            }
        }

        return result;
    }

}

Kruskal算法：

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;

public class KruskalMST {

    public static List<Edge> kruskalMST(Graph graph) {
        List<Edge> result = new ArrayList<>();
        UnionFind uf = new UnionFind(graph.vertices);

        List<Edge> edges = new ArrayList<>();
        for (int i = 0; i < graph.vertices; i++) {
            edges.addAll(graph.edges[i]);
        }

        // sort edges by weight
        Collections.sort(edges, Comparator.comparingInt(a -> a.weight));

        for (Edge e : edges) {
            int root1 = uf.find(e.from);
            int root2 = uf.find(e.to);

            if (root1 != root2) {
                uf.union(root1, root2);
                result.add(e);
            }
        }

        return result;
    }

}

测试代码：

import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;

import java.util.List;

public class TestMST {

    @Test
    public void testPrim() {
        Graph g = new Graph(5);
        g.addEdge(0, 1, 2);
        g.addEdge(0, 3, 6);
        g.addEdge(1, 3, 8);
        g.addEdge(1, 2, 3);
        g.addEdge(1, 4, 5);
        g.addEdge(2, 4, 7);
        g.addEdge(3, 4, 9);

        List<Edge> result = PrimMST.primMST(g);

        Assertions.assertEquals(result.size(), 4);
        Assertions.assertTrue(result.contains(new Edge(0, 1, 2)));
        Assertions.assertTrue(result.contains(new Edge(1, 2, 3)));
        Assertions.assertTrue(result.contains(new Edge(1, 4, 5)));
        Assertions.assertTrue(result.contains(new Edge(0, 3, 6)));
    }

    @Test
    public void testKruskal() {
        Graph g = new Graph(5);
        g.addEdge(0, 1, 2);
        g.addEdge(0, 3, 6);
        g.addEdge(1, 3, 8);
        g.addEdge(1, 2, 3);
        g.addEdge(1, 4, 5);
        g.addEdge(2, 4, 7);
        g.addEdge(3, 4, 9);

        List<Edge> result = KruskalMST.kruskalMST(g);

        Assertions.assertEquals(result.size(), 4);
        Assertions.assertTrue(result.contains(new Edge(0, 1, 2)));
        Assertions.assertTrue(result.contains(new Edge(1, 2, 3)));
        Assertions.assertTrue(result.contains(new Edge(1, 4, 5)));
        Assertions.assertTrue(result.contains(new Edge(0, 3, 6)));
    }

}