add a solution for 378

fishercoder1534 · fishercoder1534 · commit 58f4bfae0570 · 2021-07-07T10:20:55.000-07:00
diff --git a/src/main/java/com/fishercoder/solutions/_378.java b/src/main/java/com/fishercoder/solutions/_378.java
@@ -3,6 +3,7 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.PriorityQueue;
 
 public class _378 {
     public static class Solution1 {
@@ -52,4 +53,24 @@ public int kthSmallest(int[][] matrix, int k) {
             return lo;
         }
     }
+
+    public static class Solution3 {
+        /**
+         * use heap data structure
+         */
+        public int kthSmallest(int[][] matrix, int k) {
+            PriorityQueue<int[]> heap = new PriorityQueue<>((a, b) -> a[0] - b[0]);
+            for (int i = 0; i < matrix.length; i++) {
+                //we store value, rowIndex, colIndex as an array into this heap
+                heap.offer(new int[]{matrix[i][0], i, 0});
+            }
+            while (k-- > 1) {
+                int[] min = heap.poll();
+                if (min[2] + 1 < matrix[min[1]].length) {
+                    heap.offer(new int[]{matrix[min[1]][min[2] + 1], min[1], min[2] + 1});
+                }
+            }
+            return heap.poll()[0];
+        }
+    }
 }
diff --git a/src/test/java/com/fishercoder/_378Test.java b/src/test/java/com/fishercoder/_378Test.java
@@ -10,12 +10,14 @@
 public class _378Test {
     private static _378.Solution1 solution1;
     private static _378.Solution2 solution2;
+    private static _378.Solution3 solution3;
     private static int[][] matrix;
 
     @BeforeClass
     public static void setup() {
         solution1 = new _378.Solution1();
         solution2 = new _378.Solution2();
+        solution3 = new _378.Solution3();
     }
 
     @Test
@@ -25,13 +27,23 @@ public void test1() {
         };
         assertEquals(-5, solution1.kthSmallest(matrix, 1));
         assertEquals(-5, solution2.kthSmallest(matrix, 1));
+        assertEquals(-5, solution3.kthSmallest(matrix, 1));
     }
 
     @Test
     public void test2() {
         matrix = CommonUtils.convertLeetCodeIrregularLengths2DArrayInputIntoJavaArray("[1,2],[1,3]");
         assertEquals(1, solution1.kthSmallest(matrix, 2));
         assertEquals(1, solution2.kthSmallest(matrix, 2));
+        assertEquals(1, solution3.kthSmallest(matrix, 2));
+    }
+
+    @Test
+    public void test3() {
+        matrix = CommonUtils.convertLeetCodeIrregularLengths2DArrayInputIntoJavaArray("[1,5,9],[10,11,13],[12,13,15]");
+        assertEquals(13, solution1.kthSmallest(matrix, 8));
+        assertEquals(13, solution2.kthSmallest(matrix, 8));
+        assertEquals(13, solution3.kthSmallest(matrix, 8));
     }
 
 }
