refactor: improving MedianOfMatrix (#6376)

alxkm · web-flow · commit 95116dbee417 · 2025-07-15T07:23:49.000+02:00
refactor: improving MedianOfMatrix
diff --git a/src/main/java/com/thealgorithms/matrix/MedianOfMatrix.java b/src/main/java/com/thealgorithms/matrix/MedianOfMatrix.java
@@ -14,19 +14,19 @@ private MedianOfMatrix() {
     }
 
     public static int median(Iterable<List<Integer>> matrix) {
-        // Flatten the matrix into a 1D list
-        List<Integer> linear = new ArrayList<>();
+        List<Integer> flattened = new ArrayList<>();
+
         for (List<Integer> row : matrix) {
-            linear.addAll(row);
+            if (row != null) {
+                flattened.addAll(row);
+            }
         }
 
-        // Sort the 1D list
-        Collections.sort(linear);
-
-        // Calculate the middle index
-        int mid = (0 + linear.size() - 1) / 2;
+        if (flattened.isEmpty()) {
+            throw new IllegalArgumentException("Matrix must contain at least one element.");
+        }
 
-        // Return the median
-        return linear.get(mid);
+        Collections.sort(flattened);
+        return flattened.get((flattened.size() - 1) / 2);
     }
 }
diff --git a/src/test/java/com/thealgorithms/matrix/MedianOfMatrixTest.java b/src/test/java/com/thealgorithms/matrix/MedianOfMatrixTest.java
@@ -1,9 +1,11 @@
 package com.thealgorithms.matrix;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 import org.junit.jupiter.api.Test;
 
@@ -31,4 +33,18 @@ public void testMedianWithEvenNumberOfElements() {
 
         assertEquals(2, result);
     }
+
+    @Test
+    public void testMedianSingleElement() {
+        List<List<Integer>> matrix = new ArrayList<>();
+        matrix.add(List.of(1));
+
+        assertEquals(1, MedianOfMatrix.median(matrix));
+    }
+
+    @Test
+    void testEmptyMatrixThrowsException() {
+        Iterable<List<Integer>> emptyMatrix = Collections.emptyList();
+        assertThrows(IllegalArgumentException.class, () -> MedianOfMatrix.median(emptyMatrix), "Expected median() to throw, but it didn't");
+    }
 }
