LeetCode 378. Kth Smallest Element in a Sorted Matrix

原题链接在这里：https://leetcode.com/problems/kth-smallest-element-in-a-sorted-matrix/

题目：

Given a n x n matrix where each of the rows and columns are sorted in ascending order, find the kth smallest element in the matrix.

Note that it is the kth smallest element in the sorted order, not the kth distinct element.

Example:

matrix = [
   [ 1,  5,  9],
   [10, 11, 13],
   [12, 13, 15]
],
k = 8,

return 13.

Note: 
You may assume k is always valid, 1 ≤ k ≤ n2.

题解：

每一行每一列都是有序的，用一个minHeap先存入第一行，然后拿出minHeap顶部元素把它下方的元素放进minHeap中, 重复k-1次后, minHeap顶的就是第k小的元素.

Time Complexity: O(n + klogn), heapify 用时O(n), (k-1)次poll 和 offer, 每次用时O(logn).

Space: O(n), minHeap大小.

AC Java:

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        PriorityQueue<int []> minHeap = new PriorityQueue<>((a, b) -> a[2] - b[2]);
        for(int j = 0; j < n; j++){
            minHeap.add(new int[]{0, j, matrix[0][j]});
        }
        
        while(--k > 0){
            int [] cur = minHeap.poll();
            if(cur[0] == n - 1){
                continue;
            }
            
            minHeap.add(new int[]{cur[0] + 1, cur[1], matrix[cur[0] + 1][cur[1]]});
        }
        
        return minHeap.peek()[2];
    }
}

若采用Binary Search, 先猜一个值，这个值是由最小值matrix[0][0]和最大值matrix[n-1][n-1]求mid得来.

然后算下小于等于这个mid的matrix元素个数，若是个数小于k, 就在mid到最大值之间再猜一个新的mid. 直到l 不再小于 r.

Time Complexity: O(nlog(matrix[n-1][n-1] - matrix[0][0])), 每次lowerThanMidCount用时O(n).

Space: O(1).

AC Java:

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

类似Find K Pairs with Smallest Sums, Search a 2D Matrix II.