Leetcode 378. Kth Smallest Element in a Sorted Matrix

O(mlogn), m is number of columns and n is number of rows.

O(m*n) space.

Merge the matrix into an array, return kth element.

public class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        if (matrix == null || matrix.length == 0) {
            return 0;
        }
        
        int[] res = divideConquer(0, matrix.length-1, matrix);
        return res[k-1];
    }
    
    private static int[] divideConquer(int start, int end, int[][] matrix) {
        if (start == end) {
            return matrix[start];
        }
        
        int mid = (end-start)/2+start;
        int[] a1 = divideConquer(start, mid, matrix);
        int[] a2 = divideConquer(mid+1, end, matrix);
        
        return merge(a1, a2);
    }
    
    private static int[] merge(int[] a1, int[] a2) {
        int[] res = new int[a1.length+a2.length];
        int i=0, j=0, k=0;
        while (i<a1.length && j<a2.length) {
            if (a1[i] < a2[j]) {
                res[k++] = a1[i++];
            } else {
                res[k++] = a2[j++];
            }
        }
        while (i<a1.length) {
            res[k++] = a1[i++];
        }
        while (j<a2.length) {
            res[k++] = a2[j++];
        }
        return res;
    }
}

Using heap. A max heap to save k smallest elements, return top of the heap after iteration. 

public class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        PriorityQueue<Integer> maxHeap = new PriorityQueue<Integer>(
            new Comparator<Integer> () {
                public int compare (Integer a, Integer b) {
                    return b - a;
                } 
            });
            
        for (int i=0; i<matrix.length; i++) {
            for (int j=0; j<matrix[0].length; j++) {
                if (maxHeap.size() < k) {
                    maxHeap.offer(matrix[i][j]);
                } else {
                    if (matrix[i][j] < maxHeap.peek()) {
                        maxHeap.poll();
                        maxHeap.offer(matrix[i][j]);
                    }
                }
            }
        }
        
        return maxHeap.peek();
    }
}