Enhance docs, remove main, add tests in PrimMST (TheAlgorithms#5969)

Hardvan · web-flow · commit 62c9309a3162 · 2024-10-26T18:10:19.000Z
diff --git a/DIRECTORY.md b/DIRECTORY.md
@@ -852,6 +852,7 @@
               * [KosarajuTest](https://github.com/TheAlgorithms/Java/blob/master/src/test/java/com/thealgorithms/datastructures/graphs/KosarajuTest.java)
               * [KruskalTest](https://github.com/TheAlgorithms/Java/blob/master/src/test/java/com/thealgorithms/datastructures/graphs/KruskalTest.java)
               * [MatrixGraphsTest](https://github.com/TheAlgorithms/Java/blob/master/src/test/java/com/thealgorithms/datastructures/graphs/MatrixGraphsTest.java)
+              * [PrimMSTTest](https://github.com/TheAlgorithms/Java/blob/master/src/test/java/com/thealgorithms/datastructures/graphs/PrimMSTTest.java)
               * [TarjansAlgorithmTest](https://github.com/TheAlgorithms/Java/blob/master/src/test/java/com/thealgorithms/datastructures/graphs/TarjansAlgorithmTest.java)
               * [WelshPowellTest](https://github.com/TheAlgorithms/Java/blob/master/src/test/java/com/thealgorithms/datastructures/graphs/WelshPowellTest.java)
             * hashmap
diff --git a/src/main/java/com/thealgorithms/datastructures/graphs/PrimMST.java b/src/main/java/com/thealgorithms/datastructures/graphs/PrimMST.java
@@ -1,19 +1,17 @@
 package com.thealgorithms.datastructures.graphs;
 
 /**
- * A Java program for Prim's Minimum Spanning Tree (MST) algorithm. adjacency
- * matrix representation of the graph
+ * A Java program for Prim's Minimum Spanning Tree (MST) algorithm.
+ * Adjacency matrix representation of the graph.
  */
-class PrimMST {
+public class PrimMST {
 
     // Number of vertices in the graph
-
     private static final int V = 5;
 
-    // A utility function to find the vertex with minimum key
-    // value, from the set of vertices not yet included in MST
+    // A utility function to find the vertex with the minimum key
+    // value, from the set of vertices not yet included in the MST
     int minKey(int[] key, Boolean[] mstSet) {
-        // Initialize min value
         int min = Integer.MAX_VALUE;
         int minIndex = -1;
 
@@ -27,84 +25,37 @@ int minKey(int[] key, Boolean[] mstSet) {
         return minIndex;
     }
 
-    // A utility function to print the constructed MST stored in
-    // parent[]
-    void printMST(int[] parent, int n, int[][] graph) {
-        System.out.println("Edge   Weight");
-        for (int i = 1; i < V; i++) {
-            System.out.println(parent[i] + " - " + i + "    " + graph[i][parent[i]]);
-        }
-    }
-
-    // Function to construct and print MST for a graph represented
-    //  using adjacency matrix representation
-    void primMST(int[][] graph) {
-        // Array to store constructed MST
-        int[] parent = new int[V];
-
-        // Key values used to pick minimum weight edge in cut
-        int[] key = new int[V];
+    // Function to construct MST for a graph using adjacency matrix representation
+    public int[] primMST(int[][] graph) {
+        int[] parent = new int[V]; // Array to store constructed MST
+        int[] key = new int[V]; // Key values to pick minimum weight edge
+        Boolean[] mstSet = new Boolean[V]; // Vertices not yet included in MST
 
-        // To represent set of vertices not yet included in MST
-        Boolean[] mstSet = new Boolean[V];
-
-        // Initialize all keys as INFINITE
+        // Initialize all keys as INFINITE and mstSet[] as false
         for (int i = 0; i < V; i++) {
             key[i] = Integer.MAX_VALUE;
             mstSet[i] = Boolean.FALSE;
         }
 
-        // Always include first 1st vertex in MST.
-        key[0] = 0; // Make key 0 so that this vertex is
-        // picked as first vertex
+        // Always include the first vertex in MST
+        key[0] = 0; // Make key 0 to pick the first vertex
         parent[0] = -1; // First node is always root of MST
 
         // The MST will have V vertices
         for (int count = 0; count < V - 1; count++) {
-            // Pick thd minimum key vertex from the set of vertices
-            // not yet included in MST
+            // Pick the minimum key vertex not yet included in MST
             int u = minKey(key, mstSet);
-
-            // Add the picked vertex to the MST Set
             mstSet[u] = Boolean.TRUE;
 
-            // Update key value and parent index of the adjacent
-            // vertices of the picked vertex. Consider only those
-            // vertices which are not yet included in MST
-            for (int v = 0; v < V; v++) // Update the key only if graph[u][v] is smaller than key[v] // mstSet[v] is
-                                        // false for vertices not yet included in MST // graph[u][v] is non zero only
-                                        // for adjacent vertices of m
-            {
+            // Update key value and parent index of adjacent vertices of the picked vertex
+            for (int v = 0; v < V; v++) {
                 if (graph[u][v] != 0 && !mstSet[v] && graph[u][v] < key[v]) {
                     parent[v] = u;
                     key[v] = graph[u][v];
                 }
             }
         }
 
-        // print the constructed MST
-        printMST(parent, V, graph);
-    }
-
-    public static void main(String[] args) {
-        /* Let us create the following graph
-       2    3
-    (0)--(1)--(2)
-    |    / \   |
-    6| 8/   \5 |7
-    | /      \ |
-    (3)-------(4)
-         9          */
-        PrimMST t = new PrimMST();
-        int[][] graph = new int[][] {
-            {0, 2, 0, 6, 0},
-            {2, 0, 3, 8, 5},
-            {0, 3, 0, 0, 7},
-            {6, 8, 0, 0, 9},
-            {0, 5, 7, 9, 0},
-        };
-
-        // Print the solution
-        t.primMST(graph);
+        return parent; // Return the MST parent array
     }
 }
diff --git a/src/test/java/com/thealgorithms/datastructures/graphs/PrimMSTTest.java b/src/test/java/com/thealgorithms/datastructures/graphs/PrimMSTTest.java
@@ -0,0 +1,54 @@
+package com.thealgorithms.datastructures.graphs;
+
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+
+import org.junit.jupiter.api.Test;
+
+public class PrimMSTTest {
+
+    private final PrimMST primMST = new PrimMST();
+
+    @Test
+    public void testSimpleGraph() {
+        // Test graph with 5 nodes and weighted edges
+        int[][] graph = {{0, 2, 0, 6, 0}, {2, 0, 3, 8, 5}, {0, 3, 0, 0, 7}, {6, 8, 0, 0, 9}, {0, 5, 7, 9, 0}};
+
+        int[] expectedParent = {-1, 0, 1, 0, 1};
+        int[] actualParent = primMST.primMST(graph);
+
+        assertArrayEquals(expectedParent, actualParent);
+    }
+
+    @Test
+    public void testDisconnectedGraph() {
+        // Test case with a disconnected graph (no valid MST)
+        int[][] graph = {{0, 1, 0, 0, 0}, {1, 0, 2, 0, 0}, {0, 2, 0, 3, 0}, {0, 0, 3, 0, 4}, {0, 0, 0, 4, 0}};
+
+        int[] expectedParent = {-1, 0, 1, 2, 3}; // Expected MST parent array
+        int[] actualParent = primMST.primMST(graph);
+
+        assertArrayEquals(expectedParent, actualParent);
+    }
+
+    @Test
+    public void testAllEqualWeightsGraph() {
+        // Test case where all edges have equal weight
+        int[][] graph = {{0, 1, 1, 1, 1}, {1, 0, 1, 1, 1}, {1, 1, 0, 1, 1}, {1, 1, 1, 0, 1}, {1, 1, 1, 1, 0}};
+
+        int[] expectedParent = {-1, 0, 0, 0, 0}; // Expected MST parent array (any valid spanning tree)
+        int[] actualParent = primMST.primMST(graph);
+
+        assertArrayEquals(expectedParent, actualParent);
+    }
+
+    @Test
+    public void testSparseGraph() {
+        // Test case with a sparse graph (few edges)
+        int[][] graph = {{0, 1, 0, 0, 0}, {1, 0, 1, 0, 0}, {0, 1, 0, 1, 0}, {0, 0, 1, 0, 1}, {0, 0, 0, 1, 0}};
+
+        int[] expectedParent = {-1, 0, 1, 2, 3}; // Expected MST parent array
+        int[] actualParent = primMST.primMST(graph);
+
+        assertArrayEquals(expectedParent, actualParent);
+    }
+}
