From c0cbd309203dc69426679f15182fff012aee1446 Mon Sep 17 00:00:00 2001
From: alxkm <alexanderklmn@gmail.com>
Date: Sun, 13 Jul 2025 10:33:45 +0200
Subject: [PATCH] refactor: improving MedianOfMatrix

---
 .../thealgorithms/matrix/MedianOfMatrix.java  | 20 +++++++++----------
 .../matrix/MedianOfMatrixTest.java            | 16 +++++++++++++++
 2 files changed, 26 insertions(+), 10 deletions(-)

diff --git a/src/main/java/com/thealgorithms/matrix/MedianOfMatrix.java b/src/main/java/com/thealgorithms/matrix/MedianOfMatrix.java
index c710c60a2d2a..1ec977af07c6 100644
--- 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
index db66bb2d187b..b9b97014f3fc 100644
--- 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");
+    }
 }