378. Kth Smallest Element in a Sorted Matrix

378. Kth Smallest Element in a Sorted Matrix

题目链接：https://leetcode.com/problems...

求矩阵里面第k小的数，首先比较容易想到的是用heap来做，maxheap或者minheap都可以，用maxheap的话把全部元素放进heap里面，同时如果heap的size大于k就弹出，保持heap的size为k，最后root的元素就是第k小的。复杂度是k + (m*n - k)logk，其中m = matrix.length, n = matrix[0].length。

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

看discussion里面有个比较有意思的heap做法：
https://discuss.leetcode.com/...
由于矩阵是有序的，我们知道每一行右边的元素总是大于左边，下面的元素总是大于上面的。所以先把第一行放进去，然后每次增加row的值，这样可以比较matrixrow和matrixrow+1哪个小，poll出小的那个。这做法和Find K Pairs with Smallest Sums异曲同工。

public class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        // heap
        PriorityQueue<int[]> minHeap = new PriorityQueue<>(k+1, (a, b) -> a[2] - b[2]);
        for(int j = 0; j < Math.min(k, matrix[0].length); j++) {
            minHeap.offer(new int[] {0, j, matrix[0][j]});
        }
        // for k loop find the result
        for(int i = 0; i < k-1; i++) {
            int[] cur = minHeap.poll();
            if(cur[0] + 1 < matrix.length) {
                minHeap.offer(new int[] {cur[0] + 1, cur[1], matrix[cur[0] + 1][cur[1]]});
            }
        }
        
        return minHeap.poll()[2];
    }
}

标签还写了binary search，如果二分index的话，每次找到midx和midy之后，之后index很难表示出来。看了网上给的解法，没有二分index，二分的是结果。每次找到一个mid，然后求比它小的元素的个数，根据个数大于k还是小于k来二分。参考：
http://bookshadow.com/weblog/...

public class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        // binary search
        int n = matrix.length;
        int l = matrix[0][0], r = matrix[n-1][n-1];
        while(l + 1 < r) {
            int mid = l + (r - l) / 2;
            int num = count(matrix, mid);
            if(num >= k) r = mid;
            else l = mid;
        }
        if(count(matrix, r) <= k - 1) return r;
        return l;
    }
    
    private int count(int[][] matrix, int target) {
        int n = matrix.length;
        int result = 0;
        for(int i = 0; i < n; i++) {
            int j = 0;
            while(j < n && matrix[i][j] < target) j++;
            result += j;
        }
        return result;
    }
}

算count的时候可以优化：

public class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        // binary search
        int n = matrix.length;
        int l = matrix[0][0], r = matrix[n-1][n-1];
        while(l + 1 < r) {
            int mid = l + (r - l) / 2;
            int num = count(matrix, mid);
            if(num >= k) r = mid;
            else l = mid;
        }
        if(count(matrix, r) <= k - 1) return r;
        return l;
    }
    
    private int count(int[][] matrix, int target) {
        int n = matrix.length;
        int result = 0;
        int i = n - 1, j = 0;
        while(i >= 0 && j < n) {
            if(matrix[i][j] < target) {
                result += i + 1;
                j++;
            }
            else i--;
        }
        return result;
    }
}